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

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" ⚠️ расхождение: |v1v2| = {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())

View File

@@ -34,6 +34,11 @@ BATCH_SIZE=64
MICRO_BATCH=64 MICRO_BATCH=64
LR=0.0001 LR=0.0001
OUTPUT="outputs" 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 "========================================"
echo " CVGL Experiments on GTA-UAV" echo " CVGL Experiments on GTA-UAV"
echo " Epochs: $EPOCHS | Batch: $BATCH_SIZE" echo " Epochs: $EPOCHS | Batch: $BATCH_SIZE | Seed: $SEED"
echo " Versions: ${VERSIONS[*]}" echo " Versions: ${VERSIONS[*]}"
echo " Experiments: ${EXP_NUMS[*]}" echo " Experiments: ${EXP_NUMS[*]}"
echo " Total runs: $TOTAL" echo " Total runs: $TOTAL"
@@ -126,7 +131,8 @@ for VERSION in "${VERSIONS[@]}"; do
--batch_size $BATCH_SIZE \ --batch_size $BATCH_SIZE \
--micro_batch_size $MICRO_BATCH \ --micro_batch_size $MICRO_BATCH \
--lr $LR \ --lr $LR \
--output_dir "${OUTPUT}/${VERSION}" \ --seed $SEED \
--output_dir "${OUTPUT}/${VERSION}/seed${SEED}" \
--inject_mona --inject_mona
echo "✅ [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM complete" echo "✅ [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM complete"
@@ -138,6 +144,6 @@ echo "========================================"
echo " All $TOTAL experiments complete!" echo " All $TOTAL experiments complete!"
echo " Results:" echo " Results:"
for VERSION in "${VERSIONS[@]}"; do for VERSION in "${VERSIONS[@]}"; do
echo " ${VERSION}: ${OUTPUT}/${VERSION}/" echo " ${VERSION}: ${OUTPUT}/${VERSION}/seed${SEED}/"
done done
echo "========================================" echo "========================================"

View File

@@ -1,15 +1,8 @@
"""DataLoader для GTA-UAV с текстовыми описаниями для ОБЕИХ сторон. """DataLoader для GTA-UAV с текстовыми описаниями для ОБЕИХ сторон.
ВАЖНОЕ ИЗМЕНЕНИЕ (симметричная фьюзия): теперь каждый сэмпл содержит
ЧЕТЫРЕ компонента вместо трёх:
drone_image — картинка с дрона
drone_tokens — токенизированный текст описания дрона
satellite_image — спутниковый снимок (галерея, цель поиска)
satellite_tokens — токенизированный текст описания СПУТНИКА (НОВОЕ!)
И дрон, и спутник имеют собственное текстовое описание (сгенерированное И дрон, и спутник имеют собственное текстовое описание (сгенерированное
одной и той же VLM по своей картинке). Модель сливает (картинка+текст) одной и той же VLM по своей картинке). Модель сливает (картинка+текст)
СИММЕТРИЧНО на обеих сторонах через два отдельных TextFusionMLP, и только СИММЕТРИЧНО на обеих сторонах через два отдельных GatedFusion, и только
затем сравнивает два слитых вектора. затем сравнивает два слитых вектора.
Задача: symmetric fusion drone↔satellite retrieval (NGCG), где обе стороны Задача: 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]: def _clean_caption_text(raw_output: str) -> list[str]:
"""Очистить output от markdown-разметки (v1) и разбить на абзацы. """Очистить output от markdown-разметки (v1) и разбить на абзацы.
V1 формат содержит markdown-заголовки (**P1 — ...**) и зонные метки V1 формат содержит markdown-заголовки (**P1 — ...**) и зонные метки
(upper-left: ...). V2 формат — чистая проза. Функция обрабатывает оба. (upper-left:, upper-right:, lower-center: ...), причём метки встречаются
Применяется одинаково и к drone-, и к satellite-описаниям — формат и в начале строки, и в середине после `;`. V2 формат — чистая проза.
разметки общий для обеих сторон. Функция обрабатывает оба и применяется одинаково к drone- и satellite-
описаниям.
""" """
text = re.sub(r'\*\*P\d[^*]*\*\*\s*', '', raw_output) text = re.sub(r'\*\*P\d[^*]*\*\*\s*', '', raw_output)
text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text) text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text)
text = re.sub( text = _ZONE_LABEL_RE.sub(' ', text)
r'^(upper|middle|lower|center)[-]?(left|center|right)?:\s*',
'', text, flags=re.MULTILINE,
)
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
cleaned = [] cleaned = []
for p in paragraphs: for p in paragraphs:
merged = " ".join(line.strip() for line in p.split("\n") if line.strip()) merged = " ".join(line.strip() for line in p.split("\n") if line.strip())
merged = re.sub(r'\s{2,}', ' ', merged).strip()
if merged: if merged:
cleaned.append(merged) cleaned.append(merged)
return cleaned return cleaned
@@ -204,7 +210,7 @@ class GTAUAVTextDataset(Dataset):
satellite_tokens, weight, metadata). satellite_tokens, weight, metadata).
Обе стороны несут пару (картинка, текст), которая сливается моделью Обе стороны несут пару (картинка, текст), которая сливается моделью
(TextFusionMLP) в единый вектор перед сравнением. (GatedFusion) в единый вектор перед сравнением.
Args: Args:
data_root: Корневая папка GTA-UAV. 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 = tokenize_dgtrs(sat_text, context_length=self.context_length, truncate=True)
satellite_tokens = satellite_tokens.squeeze(0) # [context_length] 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 { return {
"drone_image": drone_image, # [3, H, W] "drone_image": drone_image, # [3, H, W]
@@ -354,8 +361,8 @@ class GTAUAVTextDataset(Dataset):
"satellite_tokens": satellite_tokens, # [248] — НОВОЕ ПОЛЕ "satellite_tokens": satellite_tokens, # [248] — НОВОЕ ПОЛЕ
"weight": torch.tensor(weight, dtype=torch.float32), "weight": torch.tensor(weight, dtype=torch.float32),
"drone_img_name": drone_img_name, "drone_img_name": drone_img_name,
"drone_lat": drone_loc[0], "drone_x": drone_loc[0],
"drone_lon": drone_loc[1], "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]), "satellite_tokens": torch.stack([b["satellite_tokens"] for b in batch]),
"weight": torch.stack([b["weight"] 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_img_name": [b["drone_img_name"] for b in batch],
"drone_lat": torch.tensor([b["drone_lat"] for b in batch]), "drone_x": torch.tensor([b["drone_x"] for b in batch]),
"drone_lon": torch.tensor([b["drone_lon"] 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, batch_size: int = 64,
num_workers: int = 4, num_workers: int = 4,
image_size: int = 384, image_size: int = 384,
) -> tuple[DataLoader, DataLoader]: build_test: bool = True,
"""Создать train и test DataLoader.""" ) -> tuple[DataLoader, DataLoader | None]:
"""Создать train и (опционально) диагональный test DataLoader.
build_test=False — не строить test-набор (оценка идёт через
multi-positive eval из gta_uav_eval, который сам грузит галерею и
описания; иначе описания читаются дважды).
"""
train_dataset = GTAUAVTextDataset( train_dataset = GTAUAVTextDataset(
data_root=data_root, data_root=data_root,
pairs_meta_file=train_meta, pairs_meta_file=train_meta,
@@ -398,16 +411,6 @@ def build_dataloaders(
use_semipos=True, 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_loader = DataLoader(
train_dataset, train_dataset,
batch_size=batch_size, batch_size=batch_size,
@@ -418,6 +421,23 @@ def build_dataloaders(
drop_last=True, 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_loader = DataLoader(
test_dataset, test_dataset,
batch_size=batch_size, batch_size=batch_size,
@@ -432,4 +452,5 @@ def build_dataloaders(
"📦 DataLoaders ready: train=%d batches, test=%d batches (bs=%d, levels=%s)", "📦 DataLoaders ready: train=%d batches, test=%d batches (bs=%d, levels=%s)",
len(train_loader), len(test_loader), batch_size, text_levels, len(train_loader), len(test_loader), batch_size, text_levels,
) )
return train_loader, test_loader

View File

@@ -254,4 +254,4 @@ def build_multipos_eval(
"n_queries": len(kept_entries), "n_queries": len(kept_entries),
"n_gallery": len(gallery_names), "n_gallery": len(gallery_names),
"n_skipped": n_skipped, "n_skipped": n_skipped,
} }

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

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

View File

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

View File

@@ -58,4 +58,4 @@ class StripNetEncoder(nn.Module):
def forward(self, images: torch.Tensor) -> torch.Tensor: def forward(self, images: torch.Tensor) -> torch.Tensor:
feat = self.backbone.forward_last_features(images) # [B, 512, H/32, W/32] feat = self.backbone.forward_last_features(images) # [B, 512, H/32, W/32]
pooled = self.pool(feat).flatten(1) # [B, 512] pooled = self.pool(feat).flatten(1) # [B, 512]
return self.projection(pooled) # [B, out_dim] return self.projection(pooled) # [B, out_dim]

179
train.py
View File

@@ -1,19 +1,10 @@
"""Цикл обучения dual-encoder модели на GTA-UAV (symmetric fusion-архитектура). """Цикл обучения 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 \\ python train.py \\
--data_root /path/to/GTA-UAV \\ --data_root /path/to/GTA-UAV \\
@@ -28,22 +19,26 @@ from __future__ import annotations
import argparse import argparse
import json import json
import logging import logging
import os
import random
import sys import sys
import time import time
from contextlib import nullcontext from contextlib import nullcontext
from pathlib import Path from pathlib import Path
import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
from torch.optim import AdamW 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)) sys.path.insert(0, str(Path(__file__).resolve().parent))
from src.models.dual_encoder import build_dual_encoder, get_trainable_params from src.models.dual_encoder import build_dual_encoder, get_trainable_params
from src.losses import InfoNCELoss 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 import build_dataloaders
from src.data.gta_uav_eval import build_multipos_eval
logging.basicConfig( logging.basicConfig(
level=logging.INFO, level=logging.INFO,
@@ -52,6 +47,20 @@ logging.basicConfig(
LOGGER = logging.getLogger("cvgl.train") 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 # 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()
@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( def evaluate(
model: nn.Module, model: nn.Module,
test_loader, eval_data: dict,
device: torch.device, device: torch.device,
amp_ctx, amp_ctx,
) -> dict[str, float]: ) -> 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() model.eval()
all_drone_emb = [] query_emb = _encode_loader(
all_satellite_emb = [] model, eval_data["drone_loader"], model.encode_drone, device, amp_ctx,
)
for batch in test_loader: gallery_emb = _encode_loader(
drone_images = batch["drone_image"].to(device) model, eval_data["gallery_loader"], model.encode_satellite, device, amp_ctx,
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],
) )
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() model.train()
return metrics return metrics, per_query
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -149,8 +182,8 @@ def train_one_epoch(
amp_ctx = nullcontext() amp_ctx = nullcontext()
total_loss = 0.0 total_loss = 0.0
total_acc_t2i = 0.0 total_acc_d2s = 0.0
total_acc_i2t = 0.0 total_acc_s2d = 0.0
n_steps = 0 n_steps = 0
optimizer.zero_grad() optimizer.zero_grad()
@@ -174,8 +207,8 @@ def train_one_epoch(
scaled_loss.backward() scaled_loss.backward()
total_loss += loss.item() total_loss += loss.item()
total_acc_t2i += loss_dict["acc_t2i"].item() total_acc_d2s += loss_dict["acc_d2s"].item()
total_acc_i2t += loss_dict["acc_i2t"].item() total_acc_s2d += loss_dict["acc_s2d"].item()
n_steps += 1 n_steps += 1
if (batch_idx + 1) % grad_accumulate_steps == 0: if (batch_idx + 1) % grad_accumulate_steps == 0:
@@ -201,8 +234,8 @@ def train_one_epoch(
"acc_s2d=%.3f | τ=%.4f", "acc_s2d=%.3f | τ=%.4f",
epoch, batch_idx + 1, len(train_loader), epoch, batch_idx + 1, len(train_loader),
loss.item(), loss.item(),
loss_dict["acc_t2i"].item(), loss_dict["acc_d2s"].item(),
loss_dict["acc_i2t"].item(), loss_dict["acc_s2d"].item(),
outputs["temperature"].item(), outputs["temperature"].item(),
) )
@@ -225,8 +258,8 @@ def train_one_epoch(
return { return {
"loss": total_loss / max(n_steps, 1), "loss": total_loss / max(n_steps, 1),
"acc_t2i": total_acc_t2i / max(n_steps, 1), "acc_d2s": total_acc_d2s / max(n_steps, 1),
"acc_i2t": total_acc_i2t / max(n_steps, 1), "acc_s2d": total_acc_s2d / max(n_steps, 1),
} }
@@ -235,6 +268,7 @@ def train_one_epoch(
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def main(args): def main(args):
set_seed(args.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LOGGER.info("🚀 Device: %s", device) LOGGER.info("🚀 Device: %s", device)
log_gpu_info(device) log_gpu_info(device)
@@ -248,7 +282,7 @@ def main(args):
with open(output_dir / "config.json", "w") as f: with open(output_dir / "config.json", "w") as f:
json.dump(config, f, indent=2, default=str) json.dump(config, f, indent=2, default=str)
train_loader, test_loader = build_dataloaders( train_loader, _ = build_dataloaders(
data_root=args.data_root, data_root=args.data_root,
descriptions_path=args.descriptions_path, descriptions_path=args.descriptions_path,
text_levels=args.text_levels, text_levels=args.text_levels,
@@ -257,6 +291,18 @@ def main(args):
batch_size=args.micro_batch_size, batch_size=args.micro_batch_size,
num_workers=args.num_workers, num_workers=args.num_workers,
image_size=args.image_size, 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( model = build_dual_encoder(
@@ -292,11 +338,26 @@ def main(args):
betas=(0.9, 0.98), 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, optimizer,
T_max=args.epochs, T_max=args.epochs - warmup_epochs,
eta_min=args.lr * 0.01, 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) criterion = InfoNCELoss(label_smoothing=args.label_smoothing)
@@ -341,21 +402,25 @@ def main(args):
scheduler.step() scheduler.step()
if epoch % args.eval_every == 0 or epoch == args.epochs: 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: else:
eval_metrics = {} eval_metrics, per_query = {}, None
elapsed = time.time() - t0 elapsed = time.time() - t0
LOGGER.info( LOGGER.info(
"📈 Epoch %d/%d (%.0fs) | loss=%.4f | " "📈 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, epoch, args.epochs, elapsed,
train_metrics["loss"], train_metrics["loss"],
eval_metrics.get("recall@1", 0), eval_metrics.get("recall@1", 0),
eval_metrics.get("recall@5", 0), eval_metrics.get("recall@5", 0),
eval_metrics.get("recall@10", 0), eval_metrics.get("recall@10", 0),
eval_metrics.get("AP", 0), eval_metrics.get("mAP", 0),
) )
if epoch == 1: if epoch == 1:
@@ -382,6 +447,11 @@ def main(args):
}, },
output_dir / "best_model.pth", 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) LOGGER.info("💾 New best model (R@1=%.4f)", recall1)
torch.save( torch.save(
@@ -434,7 +504,7 @@ def parse_args():
p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--lr", type=float, default=1e-4)
p.add_argument("--weight_decay", type=float, default=0.01) p.add_argument("--weight_decay", type=float, default=0.01)
p.add_argument("--max_grad_norm", type=float, default=1.0) 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) p.add_argument("--eval_every", type=int, default=1)
# Performance # Performance
@@ -447,9 +517,12 @@ def parse_args():
# Resume / output # Resume / output
p.add_argument("--resume", action="store_true", default=False) p.add_argument("--resume", action="store_true", default=False)
p.add_argument("--output_dir", type=str, default="outputs") 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() return p.parse_args()
if __name__ == "__main__": if __name__ == "__main__":
main(parse_args()) main(parse_args())