From 05355c6a7bfadcb88ccca8819cb0ff3081454bb6 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 7 Jul 2026 18:47:38 +0300 Subject: [PATCH] Full fix --- check_coverage.py | 125 +++++++++++++++++ scripts/run_experiments.sh | 12 +- src/data/gta_uav.py | 89 +++++++----- src/data/gta_uav_eval.py | 2 +- src/losses.py | 128 ++++++++--------- src/metrics.py | 245 +++++++++++++++++++-------------- src/models/dual_encoder.py | 154 ++++++++++----------- src/models/stripnet_encoder.py | 2 +- train.py | 179 +++++++++++++++++------- 9 files changed, 599 insertions(+), 337 deletions(-) create mode 100644 check_coverage.py diff --git a/check_coverage.py b/check_coverage.py new file mode 100644 index 0000000..011f129 --- /dev/null +++ b/check_coverage.py @@ -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()) diff --git a/scripts/run_experiments.sh b/scripts/run_experiments.sh index 0e0b112..c04e353 100644 --- a/scripts/run_experiments.sh +++ b/scripts/run_experiments.sh @@ -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 "========================================" \ No newline at end of file diff --git a/src/data/gta_uav.py b/src/data/gta_uav.py index ae7d930..e37fad1 100644 --- a/src/data/gta_uav.py +++ b/src/data/gta_uav.py @@ -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, @@ -432,4 +452,5 @@ def build_dataloaders( "📦 DataLoaders ready: train=%d batches, test=%d batches (bs=%d, levels=%s)", len(train_loader), len(test_loader), batch_size, text_levels, ) - \ No newline at end of file + + return train_loader, test_loader diff --git a/src/data/gta_uav_eval.py b/src/data/gta_uav_eval.py index 095a87e..cb6ab7d 100644 --- a/src/data/gta_uav_eval.py +++ b/src/data/gta_uav_eval.py @@ -254,4 +254,4 @@ def build_multipos_eval( "n_queries": len(kept_entries), "n_gallery": len(gallery_names), "n_skipped": n_skipped, - } \ No newline at end of file + } diff --git a/src/losses.py b/src/losses.py index 9b73fe1..76d14b9 100644 --- a/src/losses.py +++ b/src/losses.py @@ -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 @@ -26,70 +19,80 @@ import torch.nn.functional as F 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, } class WeightedInfoNCELoss(nn.Module): """Weighted InfoNCE из Game4Loc для GTA-UAV. - + В GTA-UAV не все пары одинаково «правильные» — есть positive (полное совпадение) и 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, - } \ No newline at end of file + "loss_d2s": loss_d2s, + "loss_s2d": loss_s2d, + "acc_d2s": acc_d2s, + "acc_s2d": acc_s2d, + } diff --git a/src/metrics.py b/src/metrics.py index f3854a9..a35f2d1 100644 --- a/src/metrics.py +++ b/src/metrics.py @@ -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,91 +21,147 @@ 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, N]. sim[i][j] = косинусное сходство текста i и изображения j. - sim = text_embeddings @ image_embeddings.T # уже L2-нормализованы → cosine sim - - # Для каждого текста ранжируем все изображения по убыванию сходства. - # ranks[i] = на какой позиции (0-indexed) стоит правильное изображение i. - # Правильное изображение для текста i — это image i (диагональ). - sorted_indices = sim.argsort(dim=1, descending=True) # [N, N] - - # Найти позицию правильного ответа (i-е изображение для i-го текста) - ranks = torch.zeros(N, dtype=torch.long, device=sim.device) + N = query_embeddings.shape[0] + M = gallery_embeddings.shape[0] + + # Матрица косинусного сходства (эмбеддинги уже L2-нормализованы) + sim = query_embeddings @ gallery_embeddings.T # [N, M] + sorted_indices = sim.argsort(dim=1, descending=True) # [N, M] + + # 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() + for i in range(N): - rank = (sorted_indices[i] == i).nonzero(as_tuple=True)[0].item() - ranks[i] = rank - - results = {} - - # Recall@K: доля запросов, где правильный ответ в топ-K + pos_set = set(positives[i]) + if not pos_set: + continue + order = sorted_np[i] # индексы галереи по убыванию сходства + + # Позиции всех 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: - 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 / 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 @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 - - # Средний Recall@K по обоим направлениям + for key, val in q2g.items(): + results[f"q2g_{key}"] = val + for key, val in g2q.items(): + results[f"g2q_{key}"] = val + + # Средний 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,35 +170,22 @@ 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()) - + return results @@ -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}") diff --git a/src/models/dual_encoder.py b/src/models/dual_encoder.py index c79be39..2a5a26b 100644 --- a/src/models/dual_encoder.py +++ b/src/models/dual_encoder.py @@ -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) diff --git a/src/models/stripnet_encoder.py b/src/models/stripnet_encoder.py index 5f22331..66db349 100644 --- a/src/models/stripnet_encoder.py +++ b/src/models/stripnet_encoder.py @@ -58,4 +58,4 @@ class StripNetEncoder(nn.Module): def forward(self, images: torch.Tensor) -> torch.Tensor: feat = self.backbone.forward_last_features(images) # [B, 512, H/32, W/32] pooled = self.pool(feat).flatten(1) # [B, 512] - return self.projection(pooled) # [B, out_dim] \ No newline at end of file + return self.projection(pooled) # [B, out_dim] diff --git a/train.py b/train.py index d83acd0..b829354 100644 --- a/train.py +++ b/train.py @@ -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,9 +517,12 @@ 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() if __name__ == "__main__": - main(parse_args()) \ No newline at end of file + main(parse_args())