iait0/cvgl_experiments
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Full fix

Browse Source
This commit is contained in:
pikaliov
2026-07-07 18:47:38 +03:00
parent 3e95f4f618
commit 05355c6a7b
9 changed files with 599 additions and 337 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
check_coverage.py Normal file
View File

@@ -0,0 +1,125 @@
"""Сравнение ПОКРЫТИЯ captions двух наборов (v1 vs v2) — отдельно от обучения.
Отвечает на один вопрос: одинаковый ли набор картинок покрыт текстом в v1 и v2?
Если да — сравнение R@1 между версиями чистое. Если нет — разница R@1 может
объясняться покрытием, а не качеством текста (см. протокол §5.4, §0.4).
Использует ТЕ ЖЕ функции чтения, что и обучение (load_text_descriptions,
combine_text_levels), поэтому «есть текст» здесь означает ровно то же, что
видит модель: непустой текст для выбранных text_levels, а не заглушку
"No description available.".
Пример:
python check_coverage.py \
--data_root /media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR \
--captions_v1 ".../GTA-UAV-LR-captions_ v1" \
--captions_v2 ".../GTA-UAV-LR-captions_v2" \
--text_levels level1 level2
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from src.data.gta_uav import combine_text_levels, load_text_descriptions
from src.data.gta_uav_eval import collect_gallery_names
def _has_real_text(
name: str,
descriptions: dict[str, dict],
text_levels: list[str],
) -> bool:
"""True, если для картинки есть непустой текст на выбранных уровнях.
Повторяет семантику обучения: пустой результат combine_text_levels →
модель получит заглушку, т.е. текста фактически нет.
"""
desc = descriptions.get(name)
if not desc:
return False
combined = combine_text_levels(desc, text_levels)
return combined.strip() != "" and combined != "No description available."
def _covered_set(
names: list[str],
descriptions: dict[str, dict],
text_levels: list[str],
) -> set[str]:
"""Подмножество имён, у которых есть реальный текст."""
return {n for n in names if _has_real_text(n, descriptions, text_levels)}
def _report_side(
title: str,
names: list[str],
desc_v1: dict[str, dict],
desc_v2: dict[str, dict],
text_levels: list[str],
) -> None:
"""Напечатать покрытие v1/v2 и пересечение для одной стороны."""
total = len(names)
cov_v1 = _covered_set(names, desc_v1, text_levels)
cov_v2 = _covered_set(names, desc_v2, text_levels)
both = cov_v1 & cov_v2
only_v1 = cov_v1 - cov_v2
only_v2 = cov_v2 - cov_v1
def pct(n: int) -> str:
return f"{n:>6d} / {total} ({100.0 * n / max(total, 1):5.1f}%)"
print(f"\n── {title} (всего уникальных: {total}) " + "─" * 20)
print(f" покрытие v1: {pct(len(cov_v1))}")
print(f" покрытие v2: {pct(len(cov_v2))}")
print(f" есть в ОБОИХ (пересеч): {pct(len(both))}")
print(f" только v1: {len(only_v1)}")
print(f" только v2: {len(only_v2)}")
if len(cov_v1) == len(cov_v2) == len(both):
print(" ✅ наборы совпадают → покрытие не влияет на сравнение")
else:
diff = abs(len(cov_v1) - len(cov_v2))
print(
f" ⚠️ расхождение: |v1−v2| = {diff} картинок; "
f"для честного сравнения оценивайте на пересечении ({len(both)})"
)
def main(args: argparse.Namespace) -> None:
data_root = Path(args.data_root)
text_levels = list(args.text_levels)
with open(data_root / args.test_meta) as f:
pairs_meta = json.load(f)
# Дрон-запросы: уникальные имена картинок дрона в тесте.
drone_names = sorted({e["drone_img_name"] for e in pairs_meta})
# Спутниковая галерея: те же уникальные тайлы, что и в eval.
gallery_names, _ = collect_gallery_names(pairs_meta)
print("=" * 64)
print(f"COVERAGE CHECK | text_levels={text_levels} | {args.test_meta}")
print("=" * 64)
for view, names in (("drone", drone_names), ("satellite", gallery_names)):
desc_v1 = load_text_descriptions(args.captions_v1, view_type=view)
desc_v2 = load_text_descriptions(args.captions_v2, view_type=view)
_report_side(
f"{view.upper()} {'queries' if view == 'drone' else 'gallery'}",
names, desc_v1, desc_v2, text_levels,
)
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description="Сравнить покрытие captions v1 vs v2")
p.add_argument("--data_root", type=str, required=True)
p.add_argument("--test_meta", type=str, default="cross-area-drone2sate-test.json")
p.add_argument("--captions_v1", type=str, required=True)
p.add_argument("--captions_v2", type=str, required=True)
p.add_argument("--text_levels", nargs="+", default=["level1"])
return p.parse_args()
if __name__ == "__main__":
main(parse_args())

12
scripts/run_experiments.sh
View File

@@ -34,6 +34,11 @@ BATCH_SIZE=64
MICRO_BATCH=64
LR=0.0001
OUTPUT="outputs"
# Seed (протокол §0.4: dev=42; для финального вывода добить победителей на
# 123 и 456). Переопределяется из окружения, дефолт остаётся 42:
# SEED=123 bash scripts/run_experiments.sh v2 2
SEED="${SEED:-42}"
# ============================================================
# Массив экспериментов: номер → уровни текста
@@ -73,7 +78,7 @@ TOTAL=$(( ${#VERSIONS[@]} * ${#EXP_NUMS[@]} ))
echo "========================================"
echo " CVGL Experiments on GTA-UAV"
echo " Epochs: $EPOCHS | Batch: $BATCH_SIZE"
echo " Epochs: $EPOCHS | Batch: $BATCH_SIZE | Seed: $SEED"
echo " Versions: ${VERSIONS[*]}"
echo " Experiments: ${EXP_NUMS[*]}"
echo " Total runs: $TOTAL"
@@ -126,7 +131,8 @@ for VERSION in "${VERSIONS[@]}"; do
--batch_size $BATCH_SIZE \
--micro_batch_size $MICRO_BATCH \
--lr $LR \
--output_dir "${OUTPUT}/${VERSION}" \
--seed $SEED \
--output_dir "${OUTPUT}/${VERSION}/seed${SEED}" \
--inject_mona
echo "✅ [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM complete"
@@ -138,6 +144,6 @@ echo "========================================"
echo " All $TOTAL experiments complete!"
echo " Results:"
for VERSION in "${VERSIONS[@]}"; do
echo " ${VERSION}: ${OUTPUT}/${VERSION}/"
echo " ${VERSION}: ${OUTPUT}/${VERSION}/seed${SEED}/"
done
echo "========================================"

87
src/data/gta_uav.py
View File

@@ -1,15 +1,8 @@
"""DataLoader для GTA-UAV с текстовыми описаниями для ОБЕИХ сторон.
ВАЖНОЕ ИЗМЕНЕНИЕ (симметричная фьюзия): теперь каждый сэмпл содержит
ЧЕТЫРЕ компонента вместо трёх:
drone_image — картинка с дрона
drone_tokens — токенизированный текст описания дрона
satellite_image — спутниковый снимок (галерея, цель поиска)
satellite_tokens — токенизированный текст описания СПУТНИКА (НОВОЕ!)
И дрон, и спутник имеют собственное текстовое описание (сгенерированное
одной и той же VLM по своей картинке). Модель сливает (картинка+текст)
СИММЕТРИЧНО на обеих сторонах через два отдельных TextFusionMLP, и только
СИММЕТРИЧНО на обеих сторонах через два отдельных GatedFusion, и только
затем сравнивает два слитых вектора.
Задача: symmetric fusion drone↔satellite retrieval (NGCG), где обе стороны
@@ -37,24 +30,37 @@ LOGGER = logging.getLogger("cvgl.data.gta_uav")
# Загрузка и очистка текстовых описаний
# ---------------------------------------------------------------------------
# Зонная метка v1: позиционное слово (upper/lower/center...) + опциональное
# горизонтальное уточнение + двоеточие. Матчится ГДЕ УГОДНО в строке (в т.ч.
# в середине после `;`), регистронезависимо. Опциональный ведущий `;`/`,`
# съедается вместе с меткой, но пробельные границы — только [ \t], чтобы НЕ
# разрушать разрыв абзацев `\n\n` (иначе level'ы слипаются).
_ZONE_LABEL_RE = re.compile(
r'[;,]?[ \t]*'
r'\b(?:upper|lower|middle|top|bottom|center|central)'
r'(?:[-‑ ]?(?:left|right|center|centre|middle))?'
r'[ \t]*:[ \t]*',
re.IGNORECASE,
)
def _clean_caption_text(raw_output: str) -> list[str]:
"""Очистить output от markdown-разметки (v1) и разбить на абзацы.
V1 формат содержит markdown-заголовки (**P1 — ...**) и зонные метки
(upper-left: ...). V2 формат — чистая проза. Функция обрабатывает оба.
Применяется одинаково и к drone-, и к satellite-описаниям — формат
разметки общий для обеих сторон.
(upper-left:, upper-right:, lower-center: ...), причём метки встречаются
и в начале строки, и в середине после `;`. V2 формат — чистая проза.
Функция обрабатывает оба и применяется одинаково к drone- и satellite-
описаниям.
"""
text = re.sub(r'\*\*P\d[^*]*\*\*\s*', '', raw_output)
text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text)
text = re.sub(
r'^(upper|middle|lower|center)[-‑]?(left|center|right)?:\s*',
'', text, flags=re.MULTILINE,
)
text = _ZONE_LABEL_RE.sub(' ', text)
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
cleaned = []
for p in paragraphs:
merged = " ".join(line.strip() for line in p.split("\n") if line.strip())
merged = re.sub(r'\s{2,}', ' ', merged).strip()
if merged:
cleaned.append(merged)
return cleaned
@@ -204,7 +210,7 @@ class GTAUAVTextDataset(Dataset):
satellite_tokens, weight, metadata).
Обе стороны несут пару (картинка, текст), которая сливается моделью
(TextFusionMLP) в единый вектор перед сравнением.
(GatedFusion) в единый вектор перед сравнением.
Args:
data_root: Корневая папка GTA-UAV.
@@ -345,7 +351,8 @@ class GTAUAVTextDataset(Dataset):
satellite_tokens = tokenize_dgtrs(sat_text, context_length=self.context_length, truncate=True)
satellite_tokens = satellite_tokens.squeeze(0) # [context_length]
drone_loc = entry.get("drone_loc_lat_lon", [0.0, 0.0])
# В meta координаты дрона называются drone_loc_x_y (игровые X/Y), а не lat/lon.
drone_loc = entry.get("drone_loc_x_y", [0.0, 0.0])
return {
"drone_image": drone_image, # [3, H, W]
@@ -354,8 +361,8 @@ class GTAUAVTextDataset(Dataset):
"satellite_tokens": satellite_tokens, # [248] — НОВОЕ ПОЛЕ
"weight": torch.tensor(weight, dtype=torch.float32),
"drone_img_name": drone_img_name,
"drone_lat": drone_loc[0],
"drone_lon": drone_loc[1],
"drone_x": drone_loc[0],
"drone_y": drone_loc[1],
}
@@ -372,8 +379,8 @@ def collate_fn(batch: list[dict]) -> dict:
"satellite_tokens": torch.stack([b["satellite_tokens"] for b in batch]),
"weight": torch.stack([b["weight"] for b in batch]),
"drone_img_name": [b["drone_img_name"] for b in batch],
"drone_lat": torch.tensor([b["drone_lat"] for b in batch]),
"drone_lon": torch.tensor([b["drone_lon"] for b in batch]),
"drone_x": torch.tensor([b["drone_x"] for b in batch]),
"drone_y": torch.tensor([b["drone_y"] for b in batch]),
}
@@ -386,8 +393,14 @@ def build_dataloaders(
batch_size: int = 64,
num_workers: int = 4,
image_size: int = 384,
) -> tuple[DataLoader, DataLoader]:
"""Создать train и test DataLoader."""
build_test: bool = True,
) -> tuple[DataLoader, DataLoader | None]:
"""Создать train и (опционально) диагональный test DataLoader.
build_test=False — не строить test-набор (оценка идёт через
multi-positive eval из gta_uav_eval, который сам грузит галерею и
описания; иначе описания читаются дважды).
"""
train_dataset = GTAUAVTextDataset(
data_root=data_root,
pairs_meta_file=train_meta,
@@ -398,16 +411,6 @@ def build_dataloaders(
use_semipos=True,
)
test_dataset = GTAUAVTextDataset(
data_root=data_root,
pairs_meta_file=test_meta,
descriptions_path=descriptions_path,
text_levels=text_levels,
image_size=image_size,
is_train=False,
use_semipos=False,
)
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
@@ -418,6 +421,23 @@ def build_dataloaders(
drop_last=True,
)
if not build_test:
LOGGER.info(
"📦 DataLoader ready: train=%d batches (bs=%d, levels=%s); test пропущен",
len(train_loader), batch_size, text_levels,
)
return train_loader, None
test_dataset = GTAUAVTextDataset(
data_root=data_root,
pairs_meta_file=test_meta,
descriptions_path=descriptions_path,
text_levels=text_levels,
image_size=image_size,
is_train=False,
use_semipos=False,
)
test_loader = DataLoader(
test_dataset,
batch_size=batch_size,
@@ -433,3 +453,4 @@ def build_dataloaders(
len(train_loader), len(test_loader), batch_size, text_levels,
)
return train_loader, test_loader

114
src/losses.py
View File

@@ -1,21 +1,14 @@
"""Contrastive loss functions для dual-encoder retrieval.
"""Contrastive loss functions для fusion dual-encoder retrieval.
InfoNCE (он же NT-Xent, CLIP loss) — стандарт для contrastive learning.
Также включён вариант weighted-InfoNCE из Game4Loc для GTA-UAV.
Как работает InfoNCE:
У тебя батч из B пар (text_i, image_i). Строим матрицу сходства B×B.
Правильная пара — на диагонали (text_0↔image_0, text_1↔image_1, ...).
Как работает InfoNCE здесь (drone-fused vs satellite):
У тебя батч из B пар (drone_i, satellite_i). Строим матрицу сходства B×B.
Правильная пара — на диагонали (drone_0↔satellite_0, drone_1↔satellite_1, ...).
Все остальные B-1 пар в строке — «негативные примеры».
Loss = CrossEntropy по строкам (text→image) + CrossEntropy по столбцам (image→text).
Это заставляет модель:
- Приближать правильные пары (высокое сходство)
- Отдалять неправильные (низкое сходство)
Чем больше B, тем больше негативов, тем сильнее сигнал.
Поэтому batch_size=64 важен.
Loss = CrossEntropy по строкам (drone→satellite) + CrossEntropy по столбцам (satellite→drone).
Это заставляет модель приближать правильные пары и отдалять неправильные.
"""
from __future__ import annotations
@@ -28,55 +21,68 @@ class InfoNCELoss(nn.Module):
"""Symmetric InfoNCE loss (как в CLIP).
Считает loss в обе стороны:
- text→image: для каждого текста, какое изображение правильное?
- image→text: для каждого изображения, какой текст правильный?
- drone→satellite (d2s): для каждого слитого запроса дрона, какой спутник правильный?
- satellite→drone (s2d): для каждого спутника, какой запрос дрона правильный?
Усредняет оба направления.
Названия acc_d2s/acc_s2d отражают направления слияния:
drone→satellite и satellite→drone.
(text-to-image / image-to-text), хотя сейчас "text" — это слитый drone-эмбеддинг,
а "image" — спутниковый эмбеддинг.
Направления взвешиваются несимметрично (протокол §6.3): основное —
q2g (drone→satellite, d2s) с весом 0.6, вспомогательное — g2q (s2d)
с весом 0.4.
Args:
label_smoothing: Сглаживание меток. 0.0 = жёсткие метки (стандарт CLIP).
Небольшое значение (0.1) может помочь при шумных описаниях.
label_smoothing: Сглаживание меток (протокол: 0.1).
w_d2s: Вес направления drone→satellite (q2g). По умолчанию 0.6.
w_s2d: Вес направления satellite→drone (g2q). По умолчанию 0.4.
"""
def __init__(self, label_smoothing: float = 0.0) -> None:
def __init__(
self,
label_smoothing: float = 0.1,
w_d2s: float = 0.6,
w_s2d: float = 0.4,
) -> None:
super().__init__()
self.label_smoothing = label_smoothing
self.w_d2s = w_d2s
self.w_s2d = w_s2d
def forward(self, logits: torch.Tensor) -> dict[str, torch.Tensor]:
"""
Args:
logits: Матрица сходства [B, B] = text_emb @ image_emb.T / τ.
logits: Матрица сходства [B, B] = drone_emb @ satellite_emb.T / τ.
Правильные пары на диагонали.
Returns:
dict:
loss: Средний loss (скаляр).
loss_t2i: Text-to-image loss.
loss_i2t: Image-to-text loss.
acc_t2i: Accuracy text→image (для мониторинга).
acc_i2t: Accuracy image→text.
loss_d2s: Drone-to-satellite loss.
loss_s2d: Satellite-to-drone loss.
acc_d2s: Accuracy drone→satellite (для мониторинга).
acc_s2d: Accuracy satellite→drone.
"""
B = logits.shape[0]
# Метки: правильная пара для i-го текста — i-е изображение
labels = torch.arange(B, device=logits.device)
# Text → Image: строки матрицы
loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
# Image → Text: столбцы матрицы (транспонируем)
loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
loss_d2s = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
loss_s2d = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
loss = (loss_t2i + loss_i2t) / 2.0
loss = self.w_d2s * loss_d2s + self.w_s2d * loss_s2d
# Accuracy для мониторинга (не для backprop)
with torch.no_grad():
acc_t2i = (logits.argmax(dim=1) == labels).float().mean()
acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean()
acc_d2s = (logits.argmax(dim=1) == labels).float().mean()
acc_s2d = (logits.T.argmax(dim=1) == labels).float().mean()
return {
"loss": loss,
"loss_t2i": loss_t2i,
"loss_i2t": loss_i2t,
"acc_t2i": acc_t2i,
"acc_i2t": acc_i2t,
"loss_d2s": loss_d2s,
"loss_s2d": loss_s2d,
"acc_d2s": acc_d2s,
"acc_s2d": acc_s2d,
}
@@ -87,9 +93,6 @@ class WeightedInfoNCELoss(nn.Module):
совпадение) и semi-positive (частичное перекрытие по IoU). Веса
отражают степень совпадения.
Вместо жёстких меток [0,0,1,0,...] используются мягкие [0.1, 0.0, 0.7, 0.2,...],
где вес пропорционален IoU перекрытия drone-satellite пар.
Args:
label_smoothing: Дополнительное сглаживание.
"""
@@ -111,33 +114,30 @@ class WeightedInfoNCELoss(nn.Module):
B = logits.shape[0]
if weights is None:
# Fallback на обычный InfoNCE
labels = torch.arange(B, device=logits.device)
loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
loss_d2s = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
loss_s2d = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
else:
# Нормализуем веса в распределение вероятностей по каждой строке
targets_t2i = F.softmax(weights, dim=1) # для text→image
targets_i2t = F.softmax(weights.T, dim=1) # для image→text
targets_d2s = F.softmax(weights, dim=1)
targets_s2d = F.softmax(weights.T, dim=1)
# KL-divergence вместо cross-entropy (для soft labels)
log_probs_t2i = F.log_softmax(logits, dim=1)
log_probs_i2t = F.log_softmax(logits.T, dim=1)
log_probs_d2s = F.log_softmax(logits, dim=1)
log_probs_s2d = F.log_softmax(logits.T, dim=1)
loss_t2i = F.kl_div(log_probs_t2i, targets_t2i, reduction="batchmean")
loss_i2t = F.kl_div(log_probs_i2t, targets_i2t, reduction="batchmean")
loss_d2s = F.kl_div(log_probs_d2s, targets_d2s, reduction="batchmean")
loss_s2d = F.kl_div(log_probs_s2d, targets_s2d, reduction="batchmean")
loss = (loss_t2i + loss_i2t) / 2.0
loss = (loss_d2s + loss_s2d) / 2.0
with torch.no_grad():
labels = torch.arange(B, device=logits.device)
acc_t2i = (logits.argmax(dim=1) == labels).float().mean()
acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean()
acc_d2s = (logits.argmax(dim=1) == labels).float().mean()
acc_s2d = (logits.T.argmax(dim=1) == labels).float().mean()
return {
"loss": loss,
"loss_t2i": loss_t2i,
"loss_i2t": loss_i2t,
"acc_t2i": acc_t2i,
"acc_i2t": acc_i2t,
"loss_d2s": loss_d2s,
"loss_s2d": loss_s2d,
"acc_d2s": acc_d2s,
"acc_s2d": acc_s2d,
}

215
src/metrics.py
View File

@@ -1,19 +1,17 @@
"""Метрики для оценки retrieval качества.
"""Метрики для оценки качества перекрёстной геолокализации (CVGL).
После обучения мы прогоняем все тексты и изображения через энкодеры,
строим матрицу сходства и для каждого запроса ранжируем кандидатов.
R@K (один positive на запрос):
R@K = (1/N) Σ_i 𝟙[rank(i) ≤ K]
где rank(i) — позиция (1-indexed) истинного соответствия.
Recall@K — основная метрика:
«Какая доля запросов нашла правильный ответ в топ-K?»
Recall@1 = 0.65 значит: в 65% случаев правильное изображение — первое в списке.
R@K (несколько positive, конвенция «hit-if-any», Sample4Geo):
достаточно, чтобы ХОТЯ БЫ ОДИН positive попал в top-K.
AP (Average Precision):
Усреднённая точность по всем позициям ранжирования.
Учитывает не только «попал ли в топ-K», но и на какой именно позиции.
AP для одного запроса — доля позитивных предсказаний среди top-K;
mAP — среднее AP по всем запросам.
Meter-level distance (специфика GTA-UAV):
Каждое изображение имеет GPS-координаты. Мы можем посчитать
расстояние в метрах между предсказанной и реальной позицией.
q2g (query-to-gallery): БПЛА → спутник (основная метрика);
g2q (gallery-to-query): спутник → БПЛА (диагностика симметричности).
"""
from __future__ import annotations
@@ -23,90 +21,146 @@ import torch
@torch.no_grad()
def compute_retrieval_metrics(
text_embeddings: torch.Tensor,
image_embeddings: torch.Tensor,
ks: list[int] = (1, 5, 10),
compute_ap: bool = True,
query_embeddings: torch.Tensor,
gallery_embeddings: torch.Tensor,
ks: tuple[int, ...] = (1, 5, 10),
positives: list[list[int]] | None = None,
compute_map: bool = True,
return_ranks: bool = False,
) -> dict[str, float]:
"""Вычислить retrieval метрики.
Предполагается что text_embeddings[i] соответствует image_embeddings[i]
(i-й текст описывает i-е изображение).
"""Вычислить метрики поиска q2g.
Args:
text_embeddings: [N, D] нормализованные текстовые эмбеддинги.
image_embeddings: [N, D] нормализованные визуальные эмбеддинги.
ks: Значения K для Recall@K.
compute_ap: Считать ли Average Precision (медленнее).
query_embeddings: [N, D] нормализованные слитые drone-эмбеддинги.
gallery_embeddings: [M, D] нормализованные слитые satellite-эмбеддинги.
ks: Значения K для R@K.
positives: Опционально — список positive-индексов галереи для каждого
запроса (для конвенции «hit-if-any» при нескольких правильных
ответах). Если None — предполагается диагональное соответствие
query[i] ↔ gallery[i] (один positive на запрос).
compute_map: Считать ли mAP.
return_ranks: Если True — добавить в результат ключ "_best_ranks":
np.ndarray [N] с 0-indexed рангом первого попавшего positive по
каждому запросу (M, если ни один не попал). Из него выводится
hit@k = (ranks < k) для любого k — нужно для bootstrap CI и
paired-теста (протокол §8.3 / §9.2).
Returns:
dict с метриками: recall@1, recall@5, recall@10, AP, ...
dict: recall@1, recall@5, recall@10, recall@1%, mean_rank,
median_rank, mAP (+ "_best_ranks", если return_ranks=True).
"""
N = text_embeddings.shape[0]
N = query_embeddings.shape[0]
M = gallery_embeddings.shape[0]
# Матрица сходства: [N, N]. sim[i][j] = косинусное сходство текста i и изображения j.
sim = text_embeddings @ image_embeddings.T # уже L2-нормализованы → cosine sim
# Матрица косинусного сходства (эмбеддинги уже L2-нормализованы)
sim = query_embeddings @ gallery_embeddings.T # [N, M]
sorted_indices = sim.argsort(dim=1, descending=True) # [N, M]
# Для каждого текста ранжируем все изображения по убыванию сходства.
# ranks[i] = на какой позиции (0-indexed) стоит правильное изображение i.
# Правильное изображение для текста i — это image i (диагональ).
sorted_indices = sim.argsort(dim=1, descending=True) # [N, N]
# positive-множество для каждого запроса
if positives is None:
# Диагональное соответствие: для запроса i правильный ответ — gallery i
positives = [[i] for i in range(N)]
# Ранг ПЕРВОГО найденного positive (hit-if-any) для каждого запроса
best_ranks = torch.full((N,), M, dtype=torch.long)
# AP по каждому запросу (для mAP)
aps = torch.zeros(N, dtype=torch.float)
sorted_np = sorted_indices.cpu().numpy()
# Найти позицию правильного ответа (i-е изображение для i-го текста)
ranks = torch.zeros(N, dtype=torch.long, device=sim.device)
for i in range(N):
rank = (sorted_indices[i] == i).nonzero(as_tuple=True)[0].item()
ranks[i] = rank
pos_set = set(positives[i])
if not pos_set:
continue
order = sorted_np[i] # индексы галереи по убыванию сходства
results = {}
# Позиции всех positive в ранжированном списке (0-indexed)
hit_positions = [r for r, idx in enumerate(order) if idx in pos_set]
if not hit_positions:
continue
# Recall@K: доля запросов, где правильный ответ в топ-K
# Ранг первого попадания (1-indexed) — для R@K по «hit-if-any»
best_ranks[i] = hit_positions[0]
# AP запроса: среднее precision@(каждое попадание)
# precision на k-м попадании = (номер попадания) / (позиция+1)
if compute_map:
ap = 0.0
for hit_num, pos in enumerate(hit_positions, start=1):
precision_at_hit = hit_num / (pos + 1)
ap += precision_at_hit
aps[i] = ap / len(pos_set)
results: dict[str, float] = {}
# R@K: доля запросов, где первый positive попал в top-K
for k in ks:
recall = (ranks < k).float().mean().item()
results[f"recall@{k}"] = recall
results[f"recall@{k}"] = (best_ranks < k).float().mean().item()
# Recall@1%: правильный ответ в топ-1% базы
k_1pct = max(1, N // 100)
results["recall@1%"] = (ranks < k_1pct).float().mean().item()
# R@1%
k_1pct = max(1, M // 100)
results["recall@1%"] = (best_ranks < k_1pct).float().mean().item()
# Mean Rank и Median Rank (для анализа — чем меньше, тем лучше)
results["mean_rank"] = ranks.float().mean().item()
results["median_rank"] = ranks.float().median().item()
# Mean / Median rank (1-indexed для читаемости)
results["mean_rank"] = (best_ranks.float() + 1).mean().item()
results["median_rank"] = (best_ranks.float() + 1).median().item()
# Average Precision
if compute_ap:
# AP = среднее по всем запросам от 1/(rank+1)
# Это упрощённый AP для случая одного правильного ответа на запрос.
ap = (1.0 / (ranks.float() + 1)).mean().item()
results["AP"] = ap
# mAP
if compute_map:
results["mAP"] = aps.mean().item()
if return_ranks:
results["_best_ranks"] = best_ranks.cpu().numpy()
return results
@torch.no_grad()
def compute_bidirectional_metrics(
text_embeddings: torch.Tensor,
image_embeddings: torch.Tensor,
ks: list[int] = (1, 5, 10),
query_embeddings: torch.Tensor,
gallery_embeddings: torch.Tensor,
ks: tuple[int, ...] = (1, 5, 10),
positives_q2g: list[list[int]] | None = None,
positives_g2q: list[list[int]] | None = None,
return_ranks: bool = False,
) -> dict[str, float]:
"""Метрики в обе стороны: text→image и image→text.
"""Метрики в обе стороны.
Returns:
dict с префиксами t2i_ и i2t_ для каждого направления.
q2g (query-to-gallery): БПЛА → спутник — основная метрика.
g2q (gallery-to-query): спутник → БПЛА — диагностика симметричности
признакового пространства.
Args:
return_ranks: Если True — добавить ключи "q2g_best_ranks" и
"g2q_best_ranks" (np.ndarray, 0-indexed ранги первого попадания)
для последующего bootstrap CI / paired-теста. Числовые метрики
остаются float, per-query массивы лежат под отдельными ключами.
"""
t2i = compute_retrieval_metrics(text_embeddings, image_embeddings, ks=ks)
# Для image→text: переворачиваем — ищем текст по изображению
i2t = compute_retrieval_metrics(image_embeddings, text_embeddings, ks=ks)
q2g = compute_retrieval_metrics(
query_embeddings, gallery_embeddings, ks=ks, positives=positives_q2g,
return_ranks=return_ranks,
)
g2q = compute_retrieval_metrics(
gallery_embeddings, query_embeddings, ks=ks, positives=positives_g2q,
return_ranks=return_ranks,
)
q2g_ranks = q2g.pop("_best_ranks", None)
g2q_ranks = g2q.pop("_best_ranks", None)
results = {}
for key, val in t2i.items():
results[f"t2i_{key}"] = val
for key, val in i2t.items():
results[f"i2t_{key}"] = val
for key, val in q2g.items():
results[f"q2g_{key}"] = val
for key, val in g2q.items():
results[f"g2q_{key}"] = val
# Средний Recall@K по обоим направлениям
# Средний R@K по обоим направлениям
for k in ks:
results[f"avg_recall@{k}"] = (t2i[f"recall@{k}"] + i2t[f"recall@{k}"]) / 2
results[f"avg_recall@{k}"] = (q2g[f"recall@{k}"] + g2q[f"recall@{k}"]) / 2
if return_ranks:
results["q2g_best_ranks"] = q2g_ranks
results["g2q_best_ranks"] = g2q_ranks
return results
@@ -116,32 +170,19 @@ def compute_meter_distance(
query_coords: np.ndarray,
gallery_coords: np.ndarray,
) -> dict[str, float]:
"""Расстояние в метрах между предсказанной и реальной позицией.
"""Расстояние в метрах между предсказанной и истинной позицией.
Специфика GTA-UAV: каждое изображение имеет GPS-координаты.
Мы берём координаты top-1 предсказания и считаем расстояние до GT.
Args:
predicted_indices: [N] индексы top-1 предсказанных изображений.
query_coords: [N, 2] GPS-координаты запросов (lat, lon).
gallery_coords: [M, 2] GPS-координаты базы изображений.
Returns:
dict: mean_distance_m, median_distance_m, recall_at_Xm.
Берёт координаты top-1 предсказанного спутникового тайла и сравнивает
с истинными координатами запроса.
"""
predicted_coords = gallery_coords[predicted_indices] # [N, 2]
# Haversine distance (приблизительно, для малых расстояний ≈ Euclidean × scale)
# Для GTA-UAV координаты в игровых юнитах, не GPS → используем Euclidean
diffs = predicted_coords - query_coords # [N, 2]
distances = np.linalg.norm(diffs, axis=1) # [N] в метрах (или юнитах)
predicted_coords = gallery_coords[predicted_indices]
diffs = predicted_coords - query_coords
distances = np.linalg.norm(diffs, axis=1)
results = {
"mean_distance": float(distances.mean()),
"median_distance": float(np.median(distances)),
}
# Recall at distance thresholds
for threshold in (5, 10, 25, 50, 100):
results[f"recall@{threshold}m"] = float((distances < threshold).mean())
@@ -152,7 +193,7 @@ def format_metrics(metrics: dict[str, float], prefix: str = "") -> str:
"""Форматировать метрики для логирования."""
lines = []
for key, val in metrics.items():
if "recall" in key or "AP" in key or "acc" in key:
if "recall" in key or "mAP" in key or "AP" in key or "acc" in key:
lines.append(f"{prefix}{key}: {val:.4f} ({val*100:.1f}%)")
elif "rank" in key:
lines.append(f"{prefix}{key}: {val:.1f}")

154
src/models/dual_encoder.py
View File

@@ -1,37 +1,12 @@
"""Dual-encoder модель для symmetric text-guided cross-view geo-localization.
ИСТОРИЯ ИЗМЕНЕНИЙ АРХИТЕКТУРЫ (по правкам от Ярослава):
v1: текст и картинка дрона кодировались раздельно и сравнивались
между собой напрямую (CLIP-style text↔image retrieval).
v2: картинка дрона и текст-описание СЛИВАЛИСЬ в один вектор
(TextFusionMLP), спутник оставался чисто визуальным.
v3 (текущая): И дрон, И спутник имеют собственное текстовое описание.
И дрон, И спутник имеют собственное текстовое описание.
Слияние происходит СИММЕТРИЧНО на обеих сторонах:
Drone (картинка + текст → слияние) ──┐
├─→ cosine similarity → InfoNCE
Satellite (картинка + текст → слияние) ┘
Зачем симметричное слияние:
У спутникового снимка тоже есть текстовое описание (та же VLM, тот же
формат), сгенерированное по картинке спутника. Раз у обеих сторон есть
пара (картинка, текст), логично использовать оба сигнала одинаковым
образом — иначе одна сторона получает преимущество (доступ к двум
модальностям), а другая остаётся обеднённой (только картинка), что не
оправдано, если данные для текста есть и там, и там.
Зачем ДВЕ отдельные TextFusionMLP (а не одна общая для обоих видов):
Визуальные признаки дрона и спутника лежат в разных доменах (наклонный
вид с дрона vs нормированный надирный спутниковый снимок), поэтому то,
как нужно сочетать "картинка + текст" в один вектор, тоже отличается.
Текстовый энкодер (DGTRS-CLIP) используется один на двоих — это одна и
та же языковая модель независимо от того, какой вид она описывает.
Зачем residual-gate в TextFusionMLP (для обеих сторон):
Картинка — "якорь" (всегда присутствует и валидна).
Текст — дополняющий сигнал переменной полезности (level1 несёт больше
сигнала чем level3, см. результаты экспериментов). Residual-gate
позволяет модели самой регулировать вклад текста, а не сливать 50/50.
"""
from __future__ import annotations
@@ -53,21 +28,27 @@ LOGGER = logging.getLogger("cvgl.dual_encoder")
# ---------------------------------------------------------------------------
# Слияние картинки и текста (TextFusionMLP) — используется для ОБЕИХ сторон
# Слияние картинки и текста (GatedFusion) — каноническая скалярная формула
# ---------------------------------------------------------------------------
class TextFusionMLP(nn.Module):
"""Слияние визуальных признаков (дрона ИЛИ спутника) с текстовым описанием.
class GatedFusion(nn.Module):
"""Каноническое gated-слияние визуальных и текстовых признаков.
Архитектура:
image_feat ──→ Linear ──┐
├─→ concat ──→ MLP ──┐
text_feat ──→ Linear ──┘ │
├─→ gate (sigmoid)
image_proj ─────────────────────────(residual)─┴─→ fused = image_proj + gate * candidate
fused = σ(α)·img_proj + (1 − σ(α))·text_proj
image_proj выступает "якорем" (residual connection): даже если текст
неинформативен (gate → 0), слитый эмбеддинг не теряет визуальный сигнал.
Формула соответствует первоисточнику (GeoRSCLIP GatedFusion) и
протоколу §3.2: `α` — ЕДИНСТВЕННЫЙ обучаемый скаляр, инициализируется
так, что σ(α) = init_gate (вес картинки на старте). σ(α) → 1 означает
"доверяем картинке", σ(α) → 0 — "доверяем тексту".
Картинка и текст приходят из РАЗНЫХ энкодеров (StripNet 1024-d,
DGTRS 768-d), поэтому сначала проецируются в общий fused_dim и
L2-нормируются — чтобы выпуклая комбинация была осмысленной
(комбинируем единичные векторы, как в общем пространстве CLIP
первоисточника).
baseline_mode=True: gate ≡ 1.0 → текст полностью игнорируется
(нужно для baseline B0 и теста разрушения текста, протокол §8).
Используется как для дрона, так и для спутника — но как ДВА РАЗНЫХ
экземпляра с независимыми весами (см. DualEncoder), поскольку визуальные
@@ -77,7 +58,8 @@ class TextFusionMLP(nn.Module):
image_dim: Размерность визуальных признаков (вход).
text_dim: Размерность текстовых признаков (вход).
fused_dim: Размерность слитого вектора (выход).
hidden_dim: Размерность скрытого слоя MLP (по умолчанию = fused_dim).
init_gate: Начальный вес картинки σ(α) ∈ (0, 1).
baseline_mode: Если True — gate ≡ 1.0 (текст выключен).
"""
def __init__(
@@ -85,56 +67,52 @@ class TextFusionMLP(nn.Module):
image_dim: int,
text_dim: int,
fused_dim: int,
hidden_dim: int | None = None,
init_gate: float = 0.7,
baseline_mode: bool = False,
) -> None:
super().__init__()
hidden_dim = hidden_dim or fused_dim
if not 0.0 < init_gate < 1.0:
raise ValueError(f"init_gate must be in (0, 1), got {init_gate}")
self.image_proj = nn.Linear(image_dim, fused_dim)
self.text_proj = nn.Linear(text_dim, fused_dim)
self.fusion_mlp = nn.Sequential(
nn.Linear(fused_dim * 2, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, fused_dim),
)
self.gate = nn.Sequential(
nn.Linear(fused_dim * 2, fused_dim),
nn.Sigmoid(),
)
# α в логит-пространстве: σ(α) = init_gate на старте.
init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
self.alpha = nn.Parameter(init_alpha)
self.baseline_mode = baseline_mode
for m in (self.image_proj, self.text_proj):
nn.init.trunc_normal_(m.weight, std=0.02)
nn.init.zeros_(m.bias)
for m in self.fusion_mlp.modules():
if isinstance(m, nn.Linear):
nn.init.trunc_normal_(m.weight, std=0.02)
nn.init.zeros_(m.bias)
# Gate инициализируем так, чтобы изначально пропускать мало текста
# (bias смещён в отрицательную сторону → sigmoid ≈ 0 на старте).
last_gate_linear = self.gate[0]
nn.init.trunc_normal_(last_gate_linear.weight, std=0.02)
nn.init.constant_(last_gate_linear.bias, -2.0)
def forward(self, image_feat: torch.Tensor, text_feat: torch.Tensor) -> torch.Tensor:
"""Слить визуальные и текстовые признаки в один вектор.
def forward(
self,
image_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
"""Слить визуальные и текстовые признаки одной скалярной gate.
Args:
image_feat: [B, image_dim] — визуальные признаки (до проекции).
text_feat: [B, text_dim] — текстовые признаки (до проекции).
text_feat: [B, text_dim] — текстовые признаки, или None.
Returns:
[B, fused_dim] — слитый вектор (картинка как якорь + gated текст).
[B, fused_dim] — слитый вектор σ(α)·img + (1 − σ(α))·text.
"""
img = self.image_proj(image_feat)
txt = self.text_proj(text_feat)
img = F.normalize(self.image_proj(image_feat), dim=-1)
if text_feat is None or self.baseline_mode:
return img
txt = F.normalize(self.text_proj(text_feat), dim=-1)
gate = torch.sigmoid(self.alpha)
return gate * img + (1.0 - gate) * txt
concat = torch.cat([img, txt], dim=-1)
candidate = self.fusion_mlp(concat)
gate = self.gate(concat)
fused = img + gate * candidate
return fused
@property
def gate_value(self) -> float:
"""Текущий вес картинки σ(α). 1.0 = текст игнорируется."""
if self.baseline_mode:
return 1.0
return torch.sigmoid(self.alpha).item()
# ---------------------------------------------------------------------------
@@ -145,7 +123,7 @@ class DualEncoder(nn.Module):
"""Symmetric dual-encoder: слитый drone-эмбеддинг vs слитый satellite-эмбеддинг.
И дрон, и спутник проходят одинаковую по структуре, но раздельную по
весам процедуру: картинка + текст → TextFusionMLP → проекция →
весам процедуру: картинка + текст → GatedFusion → проекция →
L2-норма. Полученные эмбеддинги сравниваются через cosine similarity.
Args:
@@ -168,6 +146,8 @@ class DualEncoder(nn.Module):
fused_dim: int = 512,
shared_dim: int = 512,
temperature_init: float = 0.07,
init_gate: float = 0.7,
baseline_mode: bool = False,
) -> None:
super().__init__()
@@ -176,20 +156,25 @@ class DualEncoder(nn.Module):
self.satellite_image_encoder = satellite_image_encoder
self.fused_dim = fused_dim
self.shared_dim = shared_dim
self.baseline_mode = baseline_mode
# --- Слияние картинки + текста — ДВА отдельных модуля ---
# Раздельные веса: визуальные домены дрона и спутника различны,
# поэтому оптимальный способ "смешивания" картинки с текстом тоже
# должен быть индивидуальным для каждой стороны.
self.drone_fusion = TextFusionMLP(
self.drone_fusion = GatedFusion(
image_dim=drone_image_encoder.out_dim,
text_dim=text_encoder.embed_dim,
fused_dim=fused_dim,
init_gate=init_gate,
baseline_mode=baseline_mode,
)
self.satellite_fusion = TextFusionMLP(
self.satellite_fusion = GatedFusion(
image_dim=satellite_image_encoder.out_dim,
text_dim=text_encoder.embed_dim,
fused_dim=fused_dim,
init_gate=init_gate,
baseline_mode=baseline_mode,
)
# --- Проекционные головы в общее пространство сравнения ---
@@ -217,7 +202,7 @@ class DualEncoder(nn.Module):
@property
def temperature(self) -> torch.Tensor:
return self.log_temperature.exp().clamp(min=0.01, max=100.0)
return self.log_temperature.exp().clamp(min=0.01, max=0.1)
def encode_drone(
self,
@@ -269,6 +254,8 @@ class DualEncoder(nn.Module):
satellite_emb: [B, shared_dim]
logits: [B, B] — drone_emb @ satellite_emb.T / temperature
temperature: скаляр
drone_gate: σ(α) дрон-fusion (вес картинки, для логов)
satellite_gate: σ(α) спутник-fusion (вес картинки, для логов)
"""
drone_emb = self.encode_drone(drone_images, drone_tokens)
satellite_emb = self.encode_satellite(satellite_images, satellite_tokens)
@@ -281,6 +268,8 @@ class DualEncoder(nn.Module):
"satellite_emb": satellite_emb,
"logits": logits,
"temperature": self.temperature,
"drone_gate": self.drone_fusion.gate_value,
"satellite_gate": self.satellite_fusion.gate_value,
}
@@ -330,6 +319,8 @@ def build_dual_encoder(
freeze_image_backbone: bool = True,
inject_mona: bool = True,
mona_bottleneck: int = 64,
init_gate: float = 0.7,
baseline_mode: bool = False,
device: str = "cpu",
) -> DualEncoder:
"""Собрать симметричный DualEncoder из чекпоинтов.
@@ -341,11 +332,12 @@ def build_dual_encoder(
Визуальные энкодеры — ДВА отдельных экземпляра StripNetEncoder, по
одному на домен (дрон / спутник), с отдельными Conv-MONA адаптерами.
Слияние — ДВА отдельных TextFusionMLP (drone_fusion, satellite_fusion),
с отдельными весами, так как оптимальное смешивание "картинка+текст"
индивидуально для каждого визуального домена.
Слияние — ДВА отдельных GatedFusion (drone_fusion, satellite_fusion),
с отдельными весами (проекции + скаляр α), так как оптимальное
смешивание "картинка+текст" индивидуально для каждого визуального домена.
Формула: σ(α)·img + (1 − σ(α))·text (первоисточник / протокол §3.2).
Обучаются: оба TextFusionMLP, обе проекционные головы, оба набора
Обучаются: оба GatedFusion, обе проекционные головы, оба набора
Conv-MONA адаптеров, temperature.
Замораживаются: текстовый энкодер целиком, backbone StripNet в обоих
визуальных энкодерах.
@@ -360,6 +352,8 @@ def build_dual_encoder(
freeze_image_backbone: Заморозить backbone StripNet (в обоих).
inject_mona: Инжектировать Conv-MONA в оба энкодера.
mona_bottleneck: Bottleneck dim для Conv-MONA.
init_gate: Начальный вес картинки σ(α) в GatedFusion.
baseline_mode: Если True — gate ≡ 1.0 (B0 без текста).
device: Устройство.
"""
# 1. Текстовый энкодер (общий, заморожен, используется для обеих сторон)
@@ -404,13 +398,15 @@ def build_dual_encoder(
last_n_stages=2,
)
# 5. Собрать DualEncoder (TextFusionMLP создаются внутри __init__)
# 5. Собрать DualEncoder (GatedFusion создаются внутри __init__)
model = DualEncoder(
text_encoder=text_encoder,
drone_image_encoder=drone_image_encoder,
satellite_image_encoder=satellite_image_encoder,
fused_dim=fused_dim,
shared_dim=shared_dim,
init_gate=init_gate,
baseline_mode=baseline_mode,
)
model = model.to(device)

177
train.py
View File

@@ -1,19 +1,10 @@
"""Цикл обучения dual-encoder модели на GTA-UAV (symmetric fusion-архитектура).
ВАЖНОЕ ИЗМЕНЕНИЕ: модель теперь принимает ЧЕТЫРЕ входа вместо трёх:
drone_images, drone_tokens → сливаются в drone_emb (TextFusionMLP)
satellite_images, satellite_tokens → сливаются в satellite_emb (TextFusionMLP)
drone_emb и satellite_emb сравниваются между собой → cosine similarity → InfoNCE
И дрон, и спутник имеют собственное текстовое описание; слияние
(картинка+текст) происходит СИММЕТРИЧНО на обеих сторонах, каждая —
со своим экземпляром TextFusionMLP (веса не общие, т.к. визуальные
со своим экземпляром GatedFusion (веса не общие, т.к. визуальные
домены различаются).
Оптимизировано под RTX 4090 (24 GB VRAM, Ada Lovelace): BF16 AMP,
micro_batch=64 по умолчанию (effective batch = micro_batch при отсутствии
gradient accumulation).
Использование:
python train.py \\
--data_root /path/to/GTA-UAV \\
@@ -28,22 +19,26 @@ from __future__ import annotations
import argparse
import json
import logging
import os
import random
import sys
import time
from contextlib import nullcontext
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
sys.path.insert(0, str(Path(__file__).resolve().parent))
from src.models.dual_encoder import build_dual_encoder, get_trainable_params
from src.losses import InfoNCELoss
from src.metrics import compute_retrieval_metrics, format_metrics
from src.metrics import compute_bidirectional_metrics
from src.data.gta_uav import build_dataloaders
from src.data.gta_uav_eval import build_multipos_eval
logging.basicConfig(
level=logging.INFO,
@@ -52,6 +47,20 @@ logging.basicConfig(
LOGGER = logging.getLogger("cvgl.train")
def set_seed(seed: int) -> None:
"""Зафиксировать все источники случайности (протокол §0.4: seed=42 dev).
Покрывает python random, numpy, torch (CPU/CUDA) и порядок сэмплирования
в DataLoader/semi-positive выборе — для воспроизводимого equal-budget.
"""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
LOGGER.info("🎲 Seed fixed: %d", seed)
# ---------------------------------------------------------------------------
# GPU info
# ---------------------------------------------------------------------------
@@ -92,40 +101,64 @@ def get_amp_context(use_bf16: bool, use_fp16: bool, device: torch.device):
# ---------------------------------------------------------------------------
@torch.no_grad()
@torch.no_grad()
def _encode_loader(model, loader, encode_fn, device, amp_ctx) -> torch.Tensor:
"""Прогнать loader через encode_fn (encode_drone/encode_satellite)."""
embs = []
for batch in loader:
images = batch["image"].to(device)
tokens = batch["tokens"].to(device)
with amp_ctx:
emb = encode_fn(images, tokens)
embs.append(emb.float().cpu())
return torch.cat(embs, dim=0)
def evaluate(
model: nn.Module,
test_loader,
eval_data: dict,
device: torch.device,
amp_ctx,
) -> dict[str, float]:
"""Прогнать test set: drone (fused) vs satellite (fused) — симметрично."""
"""Multi-positive retrieval eval (протокол §6.2).
Дрон-запросы и УНИКАЛЬНАЯ спутниковая галерея кодируются раздельно;
для каждого дрона учитываются ВСЕ его positive/semi-positive тайлы
(hit-if-any). Считаются оба направления: q2g (primary) и g2q.
Args:
eval_data: dict из build_multipos_eval (drone_loader, gallery_loader,
positives_q2g, positives_g2q).
"""
model.eval()
all_drone_emb = []
all_satellite_emb = []
for batch in test_loader:
drone_images = batch["drone_image"].to(device)
drone_tokens = batch["drone_tokens"].to(device)
satellite_images = batch["satellite_image"].to(device)
satellite_tokens = batch["satellite_tokens"].to(device)
with amp_ctx:
drone_emb = model.encode_drone(drone_images, drone_tokens)
satellite_emb = model.encode_satellite(satellite_images, satellite_tokens)
all_drone_emb.append(drone_emb.float().cpu())
all_satellite_emb.append(satellite_emb.float().cpu())
all_drone_emb = torch.cat(all_drone_emb, dim=0)
all_satellite_emb = torch.cat(all_satellite_emb, dim=0)
metrics = compute_retrieval_metrics(
all_drone_emb, all_satellite_emb, ks=[1, 5, 10],
query_emb = _encode_loader(
model, eval_data["drone_loader"], model.encode_drone, device, amp_ctx,
)
gallery_emb = _encode_loader(
model, eval_data["gallery_loader"], model.encode_satellite, device, amp_ctx,
)
metrics = compute_bidirectional_metrics(
query_emb, gallery_emb, ks=(1, 5, 10),
positives_q2g=eval_data["positives_q2g"],
positives_g2q=eval_data["positives_g2q"],
return_ranks=True,
)
# Отделяем per-query массивы (0-indexed ранги первого попадания) от
# числовых метрик: они пойдут в .npy для bootstrap CI / paired-теста,
# а не в history.json (иначе JSON-дамп упадёт на np.ndarray).
per_query = {
"q2g_best_ranks": metrics.pop("q2g_best_ranks"),
"g2q_best_ranks": metrics.pop("g2q_best_ranks"),
}
# Primary-алиасы: recall@k / mAP = q2g (для best-model и collect_results).
for k in (1, 5, 10):
metrics[f"recall@{k}"] = metrics[f"q2g_recall@{k}"]
metrics["mAP"] = metrics["q2g_mAP"]
model.train()
return metrics
return metrics, per_query
# ---------------------------------------------------------------------------
@@ -149,8 +182,8 @@ def train_one_epoch(
amp_ctx = nullcontext()
total_loss = 0.0
total_acc_t2i = 0.0
total_acc_i2t = 0.0
total_acc_d2s = 0.0
total_acc_s2d = 0.0
n_steps = 0
optimizer.zero_grad()
@@ -174,8 +207,8 @@ def train_one_epoch(
scaled_loss.backward()
total_loss += loss.item()
total_acc_t2i += loss_dict["acc_t2i"].item()
total_acc_i2t += loss_dict["acc_i2t"].item()
total_acc_d2s += loss_dict["acc_d2s"].item()
total_acc_s2d += loss_dict["acc_s2d"].item()
n_steps += 1
if (batch_idx + 1) % grad_accumulate_steps == 0:
@@ -201,8 +234,8 @@ def train_one_epoch(
"acc_s2d=%.3f | τ=%.4f",
epoch, batch_idx + 1, len(train_loader),
loss.item(),
loss_dict["acc_t2i"].item(),
loss_dict["acc_i2t"].item(),
loss_dict["acc_d2s"].item(),
loss_dict["acc_s2d"].item(),
outputs["temperature"].item(),
)
@@ -225,8 +258,8 @@ def train_one_epoch(
return {
"loss": total_loss / max(n_steps, 1),
"acc_t2i": total_acc_t2i / max(n_steps, 1),
"acc_i2t": total_acc_i2t / max(n_steps, 1),
"acc_d2s": total_acc_d2s / max(n_steps, 1),
"acc_s2d": total_acc_s2d / max(n_steps, 1),
}
@@ -235,6 +268,7 @@ def train_one_epoch(
# ---------------------------------------------------------------------------
def main(args):
set_seed(args.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LOGGER.info("🚀 Device: %s", device)
log_gpu_info(device)
@@ -248,7 +282,7 @@ def main(args):
with open(output_dir / "config.json", "w") as f:
json.dump(config, f, indent=2, default=str)
train_loader, test_loader = build_dataloaders(
train_loader, _ = build_dataloaders(
data_root=args.data_root,
descriptions_path=args.descriptions_path,
text_levels=args.text_levels,
@@ -257,6 +291,18 @@ def main(args):
batch_size=args.micro_batch_size,
num_workers=args.num_workers,
image_size=args.image_size,
build_test=False, # оценка идёт через multi-positive eval ниже
)
# Multi-positive eval (протокол §6.2): уникальная галерея + positive-карты.
eval_data = build_multipos_eval(
data_root=args.data_root,
test_meta=args.test_meta,
descriptions_path=args.descriptions_path,
text_levels=args.text_levels,
image_size=args.image_size,
batch_size=args.micro_batch_size,
num_workers=args.num_workers,
)
model = build_dual_encoder(
@@ -292,11 +338,26 @@ def main(args):
betas=(0.9, 0.98),
)
scheduler = CosineAnnealingLR(
# Linear-warmup + cosine (протокол §8.1 / §0.1). Warmup зажимается до
# epochs-1, чтобы на коротких прогонах (смок epochs=1) деградировать
# к чистому косинусу без ошибок.
warmup_epochs = max(0, min(args.warmup_epochs, args.epochs - 1))
cosine = CosineAnnealingLR(
optimizer,
T_max=args.epochs,
T_max=args.epochs - warmup_epochs,
eta_min=args.lr * 0.01,
)
if warmup_epochs > 0:
warmup = LinearLR(
optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_epochs,
)
scheduler = SequentialLR(
optimizer, schedulers=[warmup, cosine], milestones=[warmup_epochs],
)
LOGGER.info("📉 LR schedule: linear-warmup(%d) + cosine(%d)", warmup_epochs, args.epochs - warmup_epochs)
else:
scheduler = cosine
LOGGER.info("📉 LR schedule: cosine (%d epochs, no warmup)", args.epochs)
criterion = InfoNCELoss(label_smoothing=args.label_smoothing)
@@ -341,21 +402,25 @@ def main(args):
scheduler.step()
if epoch % args.eval_every == 0 or epoch == args.epochs:
eval_metrics = evaluate(model, test_loader, device, amp_ctx)
eval_metrics, per_query = evaluate(model, eval_data, device, amp_ctx)
# Per-query ранги последнего eval (перезапись) — источник для
# bootstrap CI / paired-теста между вариантами (протокол §8.3).
np.save(output_dir / "last_ranks_q2g.npy", per_query["q2g_best_ranks"])
np.save(output_dir / "last_ranks_g2q.npy", per_query["g2q_best_ranks"])
else:
eval_metrics = {}
eval_metrics, per_query = {}, None
elapsed = time.time() - t0
LOGGER.info(
"📈 Epoch %d/%d (%.0fs) | loss=%.4f | "
"R@1=%.3f R@5=%.3f R@10=%.3f | AP=%.3f",
"R@1=%.3f R@5=%.3f R@10=%.3f | mAP=%.3f",
epoch, args.epochs, elapsed,
train_metrics["loss"],
eval_metrics.get("recall@1", 0),
eval_metrics.get("recall@5", 0),
eval_metrics.get("recall@10", 0),
eval_metrics.get("AP", 0),
eval_metrics.get("mAP", 0),
)
if epoch == 1:
@@ -382,6 +447,11 @@ def main(args):
},
output_dir / "best_model.pth",
)
# Ранги best-эпохи — для честного сравнения вариантов по лучшей
# точке (bootstrap CI / McNemar в compare-скрипте после обучения).
if per_query is not None:
np.save(output_dir / "best_ranks_q2g.npy", per_query["q2g_best_ranks"])
np.save(output_dir / "best_ranks_g2q.npy", per_query["g2q_best_ranks"])
LOGGER.info("💾 New best model (R@1=%.4f)", recall1)
torch.save(
@@ -434,7 +504,7 @@ def parse_args():
p.add_argument("--lr", type=float, default=1e-4)
p.add_argument("--weight_decay", type=float, default=0.01)
p.add_argument("--max_grad_norm", type=float, default=1.0)
p.add_argument("--label_smoothing", type=float, default=0.0)
p.add_argument("--label_smoothing", type=float, default=0.1)
p.add_argument("--eval_every", type=int, default=1)
# Performance
@@ -447,6 +517,9 @@ def parse_args():
# Resume / output
p.add_argument("--resume", action="store_true", default=False)
p.add_argument("--output_dir", type=str, default="outputs")
p.add_argument("--seed", type=int, default=42)
p.add_argument("--warmup_epochs", type=int, default=3,
help="Эпох линейного warmup перед косинусом (§8.1).")
return p.parse_args()