From 89cb8ab0f7fc04e55ddce832d515e96c70f89a40 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Mon, 4 May 2026 11:20:14 +0300 Subject: [PATCH] claude_refactor_v3: New train_gtauav.py, added entry point main.py, added utils --- belka_refactor_04_05_log.md | 291 +++++++ src/losses/multi_infonce.py | 6 +- src/losses/weighted_infonce.py | 5 +- src/main.py | 54 ++ src/training/train_gtauav.py | 945 ++++++++++---------- src/training/train_gtauav_old.py | 1395 ++++++++++++++++++++++++++++++ src/utils/__init__.py | 2 + src/utils/io_utils.py | 51 ++ src/utils/path_utils.py | 32 + src/utils/seed_utils.py | 23 + 10 files changed, 2299 insertions(+), 505 deletions(-) create mode 100644 belka_refactor_04_05_log.md create mode 100644 src/main.py create mode 100644 src/training/train_gtauav_old.py create mode 100644 src/utils/__init__.py create mode 100644 src/utils/io_utils.py create mode 100644 src/utils/path_utils.py create mode 100644 src/utils/seed_utils.py diff --git a/belka_refactor_04_05_log.md b/belka_refactor_04_05_log.md new file mode 100644 index 0000000..aca5ddb --- /dev/null +++ b/belka_refactor_04_05_log.md @@ -0,0 +1,291 @@ +# Шаг 4а — Что изменилось + +--- + +## 1. Новая точка входа — `src/main.py` + +Запуск тренировки переехал из `src/training/train_gtauav.py::main()` в отдельный модуль `src/main.py`. + +**Старый запуск:** +```bash +python src/training/train_gtauav.py --config conf/gtauav_balanced.gin +``` + +**Новый запуск:** +```bash +python -m src.main gtauav_balanced +``` + +`src/main.py`: +- Читает имя пресета из `sys.argv[1]` (один позиционный аргумент) +- Резолвит корень проекта через `get_proj_dir()` (поиск по маркерам `pyproject.toml`/`.git`/`in/`) +- Формирует `path2cfg = f"{proj_dir}in/config_files/"` буквально по REQUIREMENTS_GIN_STYLE.md §5 +- Вызывает `load_all_configs(path2cfg, preset_name)` — двухпроходная загрузка из `_common`-файлов и пресет-директории +- Передаёт 6 объектов конфига в `train(...)` именованными аргументами + +Никакого `argparse`, никаких CLI-overrides — все параметры в `.gin`-файлах. + +--- + +## 2. Изменённый `src/training/train_gtauav.py` + +### 2.1 — Удалено +- `import argparse`, `import gin`, `from dataclasses import dataclass, field` +- Класс `TrainConfigGTAUAV` (`@dataclass + @gin.configurable`) — все его поля переехали в 6 классов в `src/conf/` +- Module-level константы `_RGB_ROOT`, `_CAPTION_ROOT`, `_TRAIN_JSON`, `_TEST_JSON`, `_DINO_WEB`, `_DINO_SAT`, `_LRSCLIP` +- Функция `main()` с argparse и CLI-overrides + +### 2.2 — Изменена сигнатура `train()` + +Было: +```python +def train(cfg: TrainConfigGTAUAV) -> None: +``` + +Стало: +```python +def train( + pipeline_cfg: PipelineConfig, + hardware_cfg: HardwareConfig, + training_cfg: TrainingConfig, + tracking_cfg: TrackingConfig, + models_common_cfg: ModelsCommonConfig, + models_cfg: DINOv3ModelsConfig | StripNetModelsConfig | SOFIAv1ModelsConfig | SOFIAv71ModelsConfig, +) -> None: +``` + +### 2.3 — Обращения `cfg.xxx` переписаны + +По карте уникальных полей: +- `cfg.train_json`, `cfg.rgb_root`, `cfg.epochs`, `cfg.output_dir`, `cfg.seed`, ... → `pipeline_cfg.*` +- `cfg.batch_size`, `cfg.grad_accum_steps`, `cfg.use_amp`, `cfg.gradient_checkpointing`, ... → `hardware_cfg.*` +- `cfg.tau_init`, `cfg.learning_rate`, `cfg.sampler_type`, `cfg.dss_*`, ... → `training_cfg.*` +- `cfg.use_wandb`, `cfg.use_tb`, `cfg.use_gradcam`, `cfg.use_profiler`, ... → `tracking_cfg.*` +- `cfg.backbone`, `cfg.baseline_mode`, `cfg.init_gate`, `cfg.lrsclip_path` → `models_common_cfg.*` +- `cfg.dino_web_path`, `cfg.shared_encoder`, `cfg.mona_*` (DINOv3-only) → `models_cfg.*` +- `cfg.stripnet_*` (StripNet-only) → `models_cfg.*` +- `cfg.sofia_preset → models_cfg.variant_label`, `cfg.sofia_d_descriptor → models_cfg.d_descriptor`, `cfg.sofia_use_text_film_*`, `cfg.sofia_mamba_*` → `models_cfg.*` +- `cfg.sofia_v1_variant → models_cfg.variant_label`, `cfg.sofia_v1_*` → `models_cfg.*` + +### 2.4 — Sofia-модели строятся напрямую из gin + +Раньше Sofia v7.1 строился через preset-фабрику + точечные overrides: + +```python +# было +preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config} +sofia_cfg = preset_map[cfg.sofia_preset]() # строит SOFIAConfig с дефолтами размера +sofia_cfg.d_descriptor = cfg.sofia_d_descriptor # потом 8 overrides +sofia_cfg.use_text_film_uav = ... +... +``` + +Теперь `SOFIAConfig(...)` собирается напрямую из всех 40+ полей `SOFIAv71ModelsConfig`: + +```python +# стало +sofia_cfg = SOFIAConfig( + input_size=models_cfg.input_size, + embed_dims=list(models_cfg.embed_dims), # все 4 dims из gin + depths=list(models_cfg.depths), + mamba_extra_kwargs=dict(models_cfg.mamba_extra_kwargs), + ... # и все остальные 35+ полей +) +``` + +Преимущество: каждый размер Sofia (Tiny/M/L) — это отдельный `presets//models.gin` со всеми полями явно. Не нужно знать, что кладёт `sofia_tiny_config()` в дефолтах. Один источник правды — gin. + +Аналогично для Sofia v1: `SOFIAv1Config(...)` строится из полей `SOFIAv1ModelsConfig`. + +### 2.5 — Direct execution убран + +```python +if __name__ == "__main__": + raise SystemExit( + "Direct execution removed. Use: python -m src.main ", + ) +``` + +--- + +## 3. Проверка передаваемых путей в энкодерах и бэкбонах + +После того как `TrainConfigGTAUAV` исчез, поля стали раскиданы по 4 семейным `Models*Config`-классам. Для того чтобы поведение **точно совпадало** со старым кодом, в каждой ветке сборки модели мы передаём **те же значения**, что приходили раньше из `cfg.*` — даже если поле теперь не имеет смысла для активного бэкбона. + +### 3.1 — Ветка `sofia_v71` + +```python +SOFIAFusionEncoder( + sofia_cfg=..., # из SOFIAv71ModelsConfig + lrsclip_path=models_common_cfg.lrsclip_path, # общий путь к DGTRS-CLIP + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + lora_rank=models_cfg.lora_rank, + device=hardware_cfg.device, +) +``` +— ничего лишнего, всё из gin. + +### 3.2 — Ветка `sofia_v1` + +```python +SOFIAv1FusionEncoder( + sofia_cfg=SOFIAv1Config(variant=models_cfg.variant_label, ...), + lrsclip_path=models_common_cfg.lrsclip_path, + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + lora_rank=models_cfg.lora_rank, + device=hardware_cfg.device, +) +``` +— симметрично с v7.1. + +### 3.3 — Ветка `stripnet` + +```python +AsymmetricEncoder( + dino_web_path="nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth", # старый _DINO_WEB + dino_sat_path="nn_models/DINO_SAT/model.safetensors", # старый _DINO_SAT + lrsclip_path=models_common_cfg.lrsclip_path, + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + shared_encoder=True, # для StripNet всегда True + mona_bottleneck=64, # старый дефолт TrainConfigGTAUAV + mona_last_n_blocks=12, # старый дефолт + device=hardware_cfg.device, + backbone=backbone, + stripnet_path=models_cfg.stripnet_path, + stripnet_mona_last_n_stages=models_cfg.stripnet_mona_last_n_stages, + stripnet_freeze=models_cfg.stripnet_freeze, +) +``` + +**Почему DINO-пути передаются для StripNet**: `AsymmetricEncoder.__init__` принимает все 13 параметров независимо от `backbone`. Для StripNet-режима DINO-пути **игнорируются** (модель строит `StripNetEncoder`, не `DINOv3ViT`). Старый код передавал те же `_DINO_WEB`/`_DINO_SAT` всегда — мы воспроизводим точно. Семантика одинакова. + +**Почему `shared_encoder=True`**: внутри `AsymmetricEncoder.__init__` на строке `if backbone == "stripnet": self.shared_encoder = True` — значение всё равно перезаписывается. Передаём `True` для семантической чистоты. + +**Почему `mona_bottleneck=64`/`mona_last_n_blocks=12`**: `mona_bottleneck=64` **используется** при `inject_conv_mona_into_stripnet(...)` для StripNet — нужно валидное значение. Старый код всегда подставлял дефолт `TrainConfigGTAUAV.mona_bottleneck=64`. Для StripNet поле `mona_bottleneck` (и `mona_last_n_blocks`, последнее не используется для StripNet) **не вынесено** в `StripNetModelsConfig` — это технический долг, отмечен ниже. Пока хардкод `64` совпадает с прежним поведением. + +### 3.4 — Ветка `dinov3` + +```python +AsymmetricEncoder( + dino_web_path=models_cfg.dino_web_path, + dino_sat_path=models_cfg.dino_sat_path, + lrsclip_path=models_common_cfg.lrsclip_path, + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + shared_encoder=models_cfg.shared_encoder, + mona_bottleneck=models_cfg.mona_bottleneck, + mona_last_n_blocks=models_cfg.mona_last_n_blocks, + device=hardware_cfg.device, + backbone=backbone, + stripnet_path="nn_models/STRIPNET/stripnet_s.pth", # старый дефолт TrainConfigGTAUAV + stripnet_mona_last_n_stages=0, + stripnet_freeze=True, +) +``` + +**Почему `stripnet_path` передаётся для DINOv3**: симметричная ситуация. `AsymmetricEncoder.__init__` принимает параметр всегда, для DINOv3 он игнорируется. Старый код передавал `cfg.stripnet_path` (дефолт `nn_models/STRIPNET/stripnet_s.pth`) даже при DINOv3 — воспроизводим то же. + +### 3.5 — Resume через `AsymmetricEncoder.load_checkpoint` + +`load_checkpoint` принимает только 4 параметра (`path`, `dino_web_path`, `dino_sat_path`, `lrsclip_path`, `device`) — остальные восстанавливаются из чекпоинта. Передаём `dino_*_path` исходя из типа `models_cfg`: +- `DINOv3ModelsConfig` → значения из конфига +- `StripNetModelsConfig` → дефолтные значения `_DINO_WEB`/`_DINO_SAT` + +Это **уже было** в старом коде; новый код просто аккуратнее распределил значения по типам конфига. + +> **Известное ограничение** (унаследовано от старого кода): `AsymmetricEncoder.load_checkpoint` не поддерживает StripNet-чекпоинты — он не принимает `backbone='stripnet'` и потому при resume StripNet-эксперимента построит DINOv3-модель. Это **не регрессия** — старый код имел тот же баг. Чинить — отдельный шаг. + +--- + +## 4. Переименование `"sofia"` → `"sofia_v71"` в if-ах + +В исходном коде разные части использовали разные имена для одного и того же бэкбона: +- В if-ах сборки модели: `if cfg.backbone == "sofia_v71"` +- В чекпоинт-блоке: `if cfg.backbone in ("sofia", "sofia_v1")` ← **остаточное старое имя** +- В сообщении gradient_checkpointing: `if cfg.backbone in ("stripnet", "sofia", "sofia_v1")` ← тоже + +Это был **остаточный баг** после промежуточного переименования `sofia → sofia_v71` — в одних местах сделали, в других забыли. На уровне runtime это не приводило к падению (sofia-эксперименты тогда не запускались), но при первом запуске sofia-пресета чекпоинт-блок не сохранил бы `sofia_cfg` для v7.1 (ветка просто не сработала бы — `backbone == "sofia_v71"` не in `("sofia", "sofia_v1")`). + +**Что сделано**: +- В **новых** `presets//models.gin` для Sofia v7.1: `ModelsCommonConfig.backbone = 'sofia_v71'` +- В новом `train_gtauav.py` **все** if-ы используют `"sofia_v71"`: + ```python + if backbone == "sofia_v71": # сборка модели + resume + enc_str + if backbone in ("sofia_v71", "sofia_v1"): # чекпоинт-блок (исправлено) + if backbone in ("stripnet", "sofia_v71", "sofia_v1"): # gradient_checkpointing (исправлено) + ``` +- В `config_loader.py` мапинг `_BACKBONE_TO_MODELS_CLS`: `"sofia_v71": SOFIAv71ModelsConfig` + +Имена теперь **согласованы** на всех уровнях: +- `src/conf/models_common_conf.py` → `backbone: str` ('dinov3' | 'stripnet' | 'sofia_v1' | 'sofia_v71') +- `src/conf/config_loader.py` → словарь маппинга +- `presets//models.gin` → биндинги +- `src/training/train_gtauav.py` → все if-ы +- `src/models/sofia_v71/` → имя директории моделей + +`"sofia"` без версии больше нигде не используется. + +--- + +## 5. Файлы, которые добавились / изменились + +### Создано +- `src/main.py` — точка входа +- `src/utils/__init__.py` +- `src/utils/path_utils.py` — `get_proj_dir()` +- `src/utils/seed_utils.py` — `set_seed()` (на 4б) +- `src/utils/io_utils.py` — `atomic_save_torch()`, `clear_vram()` (на 4б) + +### Перезаписано +- `src/training/train_gtauav.py` — новый файл (см. `train_gtauav.py` в outputs) +- `src/losses/multi_infonce.py` — снято `@gin.configurable` с `InfoNCELoss` (две строки удалены) +- `src/losses/weighted_infonce.py` — снято `@gin.configurable` с `WeightedInfoNCELoss` (две строки удалены) + +### Удалено +- `conf/` (директория, 17 файлов) — старые `.gin` мёртвый код, биндят несуществующий `TrainConfigGTAUAV` + +### Не тронуто на 4а +- `src/datasets/gtauav_dataset.py` — `_RGB_ROOT`/`_CAPTION_ROOT` остаются на module-level, но никто их теперь не использует. Удалить можно отдельным мини-коммитом или в 4б. +- `src/datasets/visloc_with_captions.py` (legacy v2) — оставлен по решению пользователя. + +--- + +## 6. Технический долг (на 4б) + +1. **`mona_bottleneck` для StripNet** — вынести из хардкода `64` в `StripNetModelsConfig.mona_bottleneck` или в `ModelsCommonConfig` +2. **Декомпозиция `train()`** на `Trainer` + методы (1100 строк → ~50 строк за метод) +3. **`_evaluate` → `src/eval/evaluator.py`** с `@torch.inference_mode()` вместо `@torch.no_grad()` +4. **`CSVLogger` → `src/training/csv_logger.py`** +5. **`_atomic_save` → `atomic_save_torch`** из `src/utils/io_utils.py` (с cleanup `.tmp` на ошибке) +6. **`_set_seed` / `_clear_vram`** заменить на `set_seed` / `clear_vram` из `src/utils/` +7. **`AsymmetricEncoder.load_checkpoint`** для StripNet — расширить сигнатуру или сделать отдельный путь resume +8. **Удалить `_RGB_ROOT`/`_CAPTION_ROOT`** из `gtauav_dataset.py` + +--- + +## 7. Контрольный smoke-test + +После применения шага 4а: + +```bash +cd + +# 1. Конфиги загружаются. +python -c " +from src.conf.config_loader import load_all_configs +from src.utils.path_utils import get_proj_dir +cfgs = load_all_configs(get_proj_dir() + 'in/config_files/', 'gtauav_balanced') +print('OK', sorted(cfgs.keys()), cfgs['models_common'].backbone) +" + +# 2. Тренировка стартует. +python -m src.main gtauav_balanced + +# 3. На 1 эпохе и 16 батчах метрики r@1_q2g/r@5_q2g/loss/tau/gate_q/gate_g +# совпадают со старым запуском до 4-го знака. +``` + +Если все три проверки проходят — шаг 4а закрыт, можно идти в 4б. \ No newline at end of file diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py index 445e9a0..c1e3a91 100644 --- a/src/losses/multi_infonce.py +++ b/src/losses/multi_infonce.py @@ -10,7 +10,6 @@ Supports both learnable temperature (CLIP-style logit_scale) and fixed/scheduled import math -import gin import torch import torch.nn as nn import torch.nn.functional as F @@ -81,10 +80,13 @@ def cosine_temperature( return tau_final + (tau_init - tau_final) * cosine -@gin.configurable class InfoNCELoss(nn.Module): """Symmetric InfoNCE with learnable or scheduled temperature. + + Note: NOT @gin.configurable. All parameters arrive explicitly from + + train() via TrainingConfig — single source of truth for gin-bindable + + values lives in src/conf/training_conf.py. + Args: temperature_init: Initial temperature value. temperature_final: Final temperature (only used if learnable=False). diff --git a/src/losses/weighted_infonce.py b/src/losses/weighted_infonce.py index 743d10d..abf05d6 100644 --- a/src/losses/weighted_infonce.py +++ b/src/losses/weighted_infonce.py @@ -14,16 +14,17 @@ WeightedInfoNCE softens this with adaptive label smoothing per sample. import math -import gin import torch import torch.nn as nn import torch.nn.functional as F -@gin.configurable class WeightedInfoNCELoss(nn.Module): """Weighted InfoNCE with adaptive per-sample label smoothing. + + Note: NOT @gin.configurable. All parameters arrive explicitly from + + train() via TrainingConfig (loss_type='weighted' branch). + For each sample i, eps_i = 1 - (1 - base_smoothing) / (1 + exp(-k * w_i)) where w_i is the positive weight (e.g. IoU with matched satellite crop). Higher weight → lower eps → sharper target (strong positive). diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..02158df --- /dev/null +++ b/src/main.py @@ -0,0 +1,54 @@ +"""Entry point: load configs and run training. + +Usage: + python -m src.main gtauav_balanced + python -m src.main gtauav_balanced_sofia_v1 +""" + +from __future__ import annotations + +import logging +import sys + +import coloredlogs + +from src.conf.config_loader import load_all_configs +from src.training.train_gtauav_old import train +from src.utils.path_utils import get_proj_dir + +logger = logging.getLogger("caption_test") + + +def main() -> None: + coloredlogs.install( + level="INFO", + logger=logger, + fmt="%(asctime)s %(name)s %(levelname)s %(message)s", + ) + + if len(sys.argv) != 2: + raise SystemExit( + "Usage: python -m src.main \n" + "Example: python -m src.main gtauav_balanced\n" + " available presets are subdirectories under in/config_files/", + ) + preset_name = sys.argv[1] + + proj_dir = get_proj_dir() + path2cfg = f"{proj_dir}in/config_files/" # per REQUIREMENTS_GIN_STYLE.md §5 + + configs = load_all_configs(path2cfg, preset_name) + + train( + pipeline_cfg=configs["pipeline"], + hardware_cfg=configs["hardware"], + training_cfg=configs["training"], + tracking_cfg=configs["tracking"], + models_common_cfg=configs["models_common"], + models_cfg=configs["models"], + ) + + +if __name__ == "__main__": + main() + diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 0655a38..a3163b8 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -5,21 +5,22 @@ from __future__ import annotations Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion. Single InfoNCE loss: query(drone+text) vs gallery(satellite). -Supports gin-config, W&B, TensorBoard, Grad-CAM, gradient monitoring, -PyTorch Profiler, and torchinfo model summary. +Supports gin-config (via src.conf), W&B, TensorBoard, Grad-CAM, gradient +monitoring, PyTorch Profiler, and torchinfo model summary. + +Note: this module no longer runs standalone. Entry point is src/main.py +(REQUIREMENTS_GIN_STYLE.md §5): + python -m src.main """ -import argparse import json import logging import math import time import warnings -from dataclasses import dataclass, field from pathlib import Path import coloredlogs -import gin import pandas as pd import torch import torch.nn as nn @@ -30,6 +31,15 @@ from torch.optim.lr_scheduler import LambdaLR from torch.utils.data import DataLoader from tqdm import tqdm +from src.conf.hardware_conf import HardwareConfig +from src.conf.models_common_conf import ModelsCommonConfig +from src.conf.models_dinov3_conf import DINOv3ModelsConfig +from src.conf.models_sofia_v1_conf import SOFIAv1ModelsConfig +from src.conf.models_sofia_v71_conf import SOFIAv71ModelsConfig +from src.conf.models_stripnet_conf import StripNetModelsConfig +from src.conf.pipeline_conf import PipelineConfig +from src.conf.tracking_conf import TrackingConfig +from src.conf.training_conf import TrainingConfig from src.datasets.gtauav_dataset import ( GTAUAVDataset, GTAUAVDroneQuery, @@ -57,127 +67,25 @@ from src.models.asymmetric_encoder import ( from src.models.sofia_fusion_encoder import SOFIAFusionEncoder from src.models.sofia_v1 import SOFIAv1Config from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder -from src.models.sofia_v71 import ( - sofia_l_config, - sofia_m_config, - sofia_tiny_config, -) +from src.models.sofia_v71 import SOFIAConfig LOGGER = logging.getLogger("caption_test.train_gtauav") -# Default paths. -_RGB_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR" -_CAPTION_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-captions" -_TRAIN_JSON = "meta/train_80.json" -_TEST_JSON = "meta/test_20.json" - -_DINO_WEB = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth" -_DINO_SAT = "nn_models/DINO_SAT/model.safetensors" -_LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt" - - -@gin.configurable(module="src.training.train_gtauav") -@dataclass -class TrainConfigGTAUAV: - """Training configuration for GTA-UAV experiment.""" - - # Data. - train_json: str = _TRAIN_JSON - test_json: str = _TEST_JSON - rgb_root: str = _RGB_ROOT - caption_root: str = _CAPTION_ROOT - filter_meta: str | None = None - - # Model. - dino_web_path: str = _DINO_WEB - dino_sat_path: str = _DINO_SAT - lrsclip_path: str = _LRSCLIP - init_gate: float = 0.7 - baseline_mode: bool = False - shared_encoder: bool = True # single DINOv3 WEB for both branches (simpler, half the params) - mona_bottleneck: int = 64 - mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks - gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch) - # StripNet backbone option (replaces DINOv3 when backbone="stripnet"). - backbone: str = "dinov3" # "dinov3", "stripnet", or "sofia" - stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth" - stripnet_mona_last_n_stages: int = 2 # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA) - stripnet_freeze: bool = True # If False, StripNet backbone is fully trainable (full fine-tune) - stripnet_backbone_lr_factor: float = 0.1 # Backbone LR = learning_rate * factor (only when unfrozen) - # SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint. - sofia_preset: str = "Tiny" # "Tiny" | "M" | "L" - sofia_d_descriptor: int = 1024 # retrieval space (1024 = match TextFusionMLP out_dim) - sofia_use_text_film_uav: bool = True # mid-level text-FiLM in UAV head - sofia_use_text_film_sat: bool = True # mid-level text-FiLM in SAT head - sofia_lora_rank: int = 4 - sofia_mamba_variant: str = "mamba2" # "mamba1" | "mamba2" | "efficient_vmamba" - sofia_mamba_backend: str = "auto" # "auto" | "torch" | "mamba_ssm" - # EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in). - sofia_use_evss_bridge: bool = False - sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"]) - # SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch. - sofia_v1_variant: str = "small" # "tiny_tiny" | "tiny" | "small" | "small_v2" - sofia_v1_dcn_variant: str = "v2" # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky) - sofia_v1_d_descriptor: int = 1024 - sofia_v1_use_text_film_uav: bool = True - sofia_v1_use_text_film_sat: bool = True - sofia_v1_use_film_altitude: bool = True - sofia_v1_lora_rank: int = 4 - - # Training. - resume_from: str | None = None # path to checkpoint for resuming - output_dir: str = "out/gtauav/with_text" - epochs: int = 10 - batch_size: int = 8 - num_workers: int = 4 - learning_rate: float = 1e-4 - text_lr_factor: float = 0.1 # text encoder LR = learning_rate * factor - weight_decay: float = 1e-4 - grad_clip: float = 1.0 - grad_accum_steps: int = 1 # gradient accumulation steps (effective_batch = batch_size * accum) - use_amp: bool = True - eval_every: int = 2 - warmup_epochs: int = 2 - seed: int = 42 - device: str = "cuda" - - # Loss. - loss_type: str = "symmetric" # "symmetric" (InfoNCE) or "weighted" (WeightedInfoNCE) - tau_init: float = 0.07 - label_smoothing: float = 0.1 - learnable_temperature: bool = True - weight_q2g: float = 0.6 - weight_g2q: float = 0.4 - neg_bank_size: int = 4096 # hard negative memory bank size (0 = disabled) - - # Sampling. - sampler_type: str = "mutex" # "mutex" (no false negatives) or "dss" (DSS + mutex) - dss_reembed_every: int = 1 # Re-embed train queries every N epochs for DSS. - dss_warmup_epochs: int = 1 # Use mutex-only for the first N epochs (fresh model embeddings aren't useful) - dss_knn_device: str = "cuda" # Device for similarity matmul in DSS sampler. - dss_use_lsh: bool = False # Approximate kNN via LSH (opt-in; exact is fast at 25K). - dss_lsh_num_tables: int = 8 - dss_lsh_num_bits: int = 14 - dss_cache_dir: str | None = None # Disk cache for embeddings; None = disabled. - # Legacy alias kept for backward compatibility. - use_mutex_sampler: bool = True - - # Tracking & diagnostics. - use_wandb: bool = False - use_tb: bool = True - wandb_project: str = "caption-test-gtauav" - wandb_run_name: str | None = None - wandb_entity: str | None = None - log_grad_norms: bool = True - use_gradcam: bool = False - gradcam_every: int = 5 # Grad-CAM every N epochs - gradcam_samples: int = 8 - use_profiler: bool = False - profiler_warmup: int = 3 - profiler_active: int = 5 +# Type alias for the family-specific models config. +ModelsConfig = ( + DINOv3ModelsConfig + | StripNetModelsConfig + | SOFIAv1ModelsConfig + | SOFIAv71ModelsConfig +) def _set_seed(seed: int) -> None: + """Fix all RNG seeds for reproducibility. + + Note: duplicates src.utils.seed_utils.set_seed. Will be removed in step 4b + when this module gets decomposed into Trainer. + """ import random as _random import numpy as _np _random.seed(seed) @@ -187,6 +95,11 @@ def _set_seed(seed: int) -> None: def _atomic_save(obj: dict, path: Path) -> None: + """Save checkpoint atomically via .tmp + replace. + + Note: will be replaced with src.utils.io_utils.atomic_save_torch in step 4b + (current version doesn't clean up .tmp on error). + """ path.parent.mkdir(parents=True, exist_ok=True) tmp_path = path.with_suffix(path.suffix + ".tmp") torch.save(obj, tmp_path) @@ -199,107 +112,106 @@ def _build_param_groups( text_lr_factor: float, stripnet_backbone_lr_factor: float = 0.1, ) -> list[dict]: - """Build optimizer param groups with separate LR for text encoder and unfrozen StripNet backbone. + """Build parameter groups with separate LR for text encoder and (optionally) StripNet backbone. - Groups: - - text_encoder.* → lr * text_lr_factor (default 1e-5) - - image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5) - - everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv, - SOFIA backbone+heads when backbone="sofia") → lr + Group 0: projections + heads + MONA + (logit_scale appended later). + Group 1: DGTRS-CLIP text encoder (lr * text_lr_factor). + Group 2 (optional): StripNet backbone when unfrozen (lr * stripnet_backbone_lr_factor). """ - text_params = [] - backbone_params = [] - other_params = [] + main_params: list[nn.Parameter] = [] + text_params: list[nn.Parameter] = [] + stripnet_backbone_params: list[nn.Parameter] = [] - is_stripnet = isinstance(getattr(model, "image_encoder", None), nn.Module) and \ - getattr(model, "backbone", "dinov3") == "stripnet" - - for name, param in model.named_parameters(): - if not param.requires_grad: + for name, p in model.named_parameters(): + if not p.requires_grad: continue if "text_encoder" in name: - text_params.append(param) - elif is_stripnet and name.startswith("image_encoder.backbone.") and "mona_" not in name: - backbone_params.append(param) + text_params.append(p) + elif name.startswith("backbone.") or name.startswith("stripnet."): + # StripNet backbone params (only present when stripnet_freeze=False). + stripnet_backbone_params.append(p) else: - other_params.append(param) + main_params.append(p) - groups = [{"params": other_params, "lr": lr}] + groups: list[dict] = [ + {"params": main_params, "lr": lr, "name": "main"}, + ] if text_params: - groups.append({"params": text_params, "lr": lr * text_lr_factor}) - if backbone_params: - groups.append({"params": backbone_params, "lr": lr * stripnet_backbone_lr_factor}) - + groups.append({"params": text_params, "lr": lr * text_lr_factor, "name": "text"}) + if stripnet_backbone_params: + groups.append({ + "params": stripnet_backbone_params, + "lr": lr * stripnet_backbone_lr_factor, + "name": "stripnet_backbone", + }) return groups -def _cosine_warmup_schedule( - warmup_steps: int, - total_steps: int, -) -> callable: - """Cosine annealing with linear warmup.""" +def _cosine_warmup_schedule(warmup_steps: int, total_steps: int): + """Return a lr_lambda for LambdaLR: linear warmup + cosine decay.""" def lr_lambda(step: int) -> float: if step < warmup_steps: - return step / max(warmup_steps, 1) - progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) - return 0.5 * (1.0 + math.cos(math.pi * progress)) + return float(step) / max(1, warmup_steps) + progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) + return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0))) return lr_lambda -@torch.no_grad() def _embed_drone_queries( - model: AsymmetricEncoder, - train_ds: GTAUAVDataset, + model: nn.Module, + train_ds, device: str, batch_size: int, num_workers: int, ) -> torch.Tensor: - """Forward all drone queries and return [N, D] embeddings on CPU. + """Embed all drone queries from train_ds using the current model state. - Used by DynamicSimilaritySampler to rank drones by visual similarity. - Runs with model.eval() but restores original train state afterwards. + Used by DSS sampler at the start of each non-warmup epoch. """ - was_training = model.training model.eval() - - query_ds = GTAUAVDroneQuery(train_ds) + drone_query_ds = GTAUAVDroneQuery( + train_ds.entries, + rgb_root=str(train_ds.rgb_root), + caption_root=str(train_ds.caption_root), + image_transform=get_dino_transform(image_size=256), + ) loader = DataLoader( - query_ds, + drone_query_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_drone_query, pin_memory=True, ) - - embs: list[torch.Tensor] = [] - for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): - drone_img = batch["drone_img"].to(device, non_blocking=True) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(device, non_blocking=True) - q = model.encode_query( - drone_img, - batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], - altitude=altitude, - ) - embs.append(q.cpu()) - - if was_training: - model.train() - return torch.cat(embs, dim=0) + all_embs: list[torch.Tensor] = [] + with torch.inference_mode(): + for batch in tqdm(loader, desc="dss-embed", unit="batch", leave=False): + drone_img = batch["drone_img"].to(device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(device, non_blocking=True) + kwargs = {"drone_img": drone_img, "altitude": altitude} + if not getattr(model, "baseline_mode", False): + kwargs["caption_l1"] = batch["caption_l1"] + kwargs["caption_l2"] = batch["caption_l2"] + kwargs["caption_l3"] = batch["caption_l3"] + with autocast(device_type="cuda", enabled=True): + emb = model.encode_drone_query(**kwargs) + all_embs.append(emb.cpu()) + model.train() + return torch.cat(all_embs, dim=0) @torch.no_grad() def _evaluate( - model: AsymmetricEncoder, + model: nn.Module, loader: DataLoader, device: str, - loss_fn: nn.Module | None = None, - epoch: int = 0, - total_epochs: int = 1, + loss_fn: nn.Module, + epoch: int, + total_epochs: int, k_values: tuple[int, ...] = (1, 5, 10), max_batches: int | None = None, desc: str = "eval", @@ -315,6 +227,7 @@ def _evaluate( `max_batches` subsamples the drone queries (not the gallery) — useful for a quick train-side sanity check. """ + dataset = loader.dataset if not isinstance(dataset, GTAUAVDataset): raise TypeError(f"_evaluate expects GTAUAVDataset, got {type(dataset).__name__}") @@ -567,143 +480,258 @@ class CSVLogger: pd.DataFrame(self.train_recall_rows).to_csv(self.log_dir / "train_recall.csv", index=False) + def _clear_vram() -> None: - """Free VRAM from previous runs before starting.""" - import gc - gc.collect() + """Free VRAM and reset peak memory stats. + + Note: duplicates src.utils.io_utils.clear_vram. Will be replaced in step 4b. + """ + import gc as _gc + _gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() - allocated = torch.cuda.memory_allocated() / 1e9 - LOGGER.info("VRAM cleared. Current usage: %.2f GB", allocated) + allocated_gb = torch.cuda.memory_allocated() / 1e9 + LOGGER.info("VRAM cleared. Current usage: %.2f GB", allocated_gb) -def train(cfg: TrainConfigGTAUAV) -> None: - """Run full training loop.""" +def train( + pipeline_cfg: PipelineConfig, + hardware_cfg: HardwareConfig, + training_cfg: TrainingConfig, + tracking_cfg: TrackingConfig, + models_common_cfg: ModelsCommonConfig, + models_cfg: ModelsConfig, +) -> None: + """Run full training loop. + + Args: + pipeline_cfg: Paths, schedule (epochs/eval_every/warmup), seed, output_dir. + hardware_cfg: batch_size, grad_accum, num_workers, AMP, gradient_checkpointing. + training_cfg: Loss + optimizer + sampler recipe. + tracking_cfg: W&B / TensorBoard / Grad-CAM / profiler. + models_common_cfg: backbone, baseline_mode, init_gate, lrsclip_path. + models_cfg: Family-specific config selected by models_common_cfg.backbone: + DINOv3ModelsConfig | StripNetModelsConfig + | SOFIAv1ModelsConfig | SOFIAv71ModelsConfig. + """ coloredlogs.install( level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s", ) _clear_vram() - _set_seed(cfg.seed) - output_dir = Path(cfg.output_dir) + _set_seed(pipeline_cfg.seed) + output_dir = Path(pipeline_cfg.output_dir) output_dir.mkdir(parents=True, exist_ok=True) - # Save config. + # Save config — all 6 config objects merged into one dict for traceability. + full_config = { + "pipeline": vars(pipeline_cfg), + "hardware": vars(hardware_cfg), + "training": vars(training_cfg), + "tracking": vars(tracking_cfg), + "models_common": vars(models_common_cfg), + "models": vars(models_cfg), + } with (output_dir / "config.json").open("w") as f: - json.dump(vars(cfg), f, indent=2) + json.dump(full_config, f, indent=2) # --- Experiment tracker (W&B + TensorBoard) --- tracker = ExperimentTracker( output_dir=output_dir, - config=vars(cfg), - use_wandb=cfg.use_wandb, - use_tb=cfg.use_tb, - wandb_project=cfg.wandb_project, - wandb_run_name=cfg.wandb_run_name, - wandb_entity=cfg.wandb_entity, + config=full_config, + use_wandb=tracking_cfg.use_wandb, + use_tb=tracking_cfg.use_tb, + wandb_project=tracking_cfg.wandb_project, + wandb_run_name=tracking_cfg.wandb_run_name, + wandb_entity=tracking_cfg.wandb_entity, ) # Model. + backbone = models_common_cfg.backbone start_epoch = 0 resume_ckpt = None - if cfg.resume_from is not None: - LOGGER.info("Resuming from %s", cfg.resume_from) - if cfg.backbone == "sofia": + if pipeline_cfg.resume_from is not None: + LOGGER.info("Resuming from %s", pipeline_cfg.resume_from) + if backbone == "sofia_v71": model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint( - cfg.resume_from, - lrsclip_path=cfg.lrsclip_path, - device=cfg.device, + pipeline_cfg.resume_from, + lrsclip_path=models_common_cfg.lrsclip_path, + device=hardware_cfg.device, ) - elif cfg.backbone == "sofia_v1": + elif backbone == "sofia_v1": model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint( - cfg.resume_from, - lrsclip_path=cfg.lrsclip_path, - device=cfg.device, + pipeline_cfg.resume_from, + lrsclip_path=models_common_cfg.lrsclip_path, + device=hardware_cfg.device, ) else: + # DINOv3 or StripNet — both go through AsymmetricEncoder. + assert isinstance(models_cfg, (DINOv3ModelsConfig, StripNetModelsConfig)), ( + f"Expected DINOv3/StripNet ModelsConfig for backbone={backbone!r}, " + f"got {type(models_cfg).__name__}" + ) + dino_web_path = ( + models_cfg.dino_web_path if isinstance(models_cfg, DINOv3ModelsConfig) + else "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth" + ) + dino_sat_path = ( + models_cfg.dino_sat_path if isinstance(models_cfg, DINOv3ModelsConfig) + else "nn_models/DINO_SAT/model.safetensors" + ) model, resume_ckpt = AsymmetricEncoder.load_checkpoint( - cfg.resume_from, - dino_web_path=cfg.dino_web_path, - dino_sat_path=cfg.dino_sat_path, - lrsclip_path=cfg.lrsclip_path, - device=cfg.device, + pipeline_cfg.resume_from, + dino_web_path=dino_web_path, + dino_sat_path=dino_sat_path, + lrsclip_path=models_common_cfg.lrsclip_path, + device=hardware_cfg.device, ) start_epoch = resume_ckpt.get("epoch", -1) + 1 else: - mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)" - if cfg.backbone == "sofia": - enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})" - elif cfg.backbone == "sofia_v1": - enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})" - elif cfg.backbone == "stripnet": + mode_str = "baseline (no text)" if models_common_cfg.baseline_mode else "with text (L1/L2/L3)" + if backbone == "sofia_v71": + assert isinstance(models_cfg, SOFIAv71ModelsConfig) + enc_str = ( + f"SOFIA-{models_cfg.variant_label} " + f"(text-FiLM uav={models_cfg.use_text_film_uav}, " + f"sat={models_cfg.use_text_film_sat})" + ) + elif backbone == "sofia_v1": + assert isinstance(models_cfg, SOFIAv1ModelsConfig) + enc_str = ( + f"SOFIAv1-{models_cfg.variant_label} (StripNet+DCNv4, " + f"text-FiLM uav={models_cfg.use_text_film_uav}, " + f"sat={models_cfg.use_text_film_sat})" + ) + elif backbone == "stripnet": enc_str = "StripNet-small (shared, 512→1024 proj)" - else: - enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)" + else: # dinov3 + assert isinstance(models_cfg, DINOv3ModelsConfig) + enc_str = "shared DINOv3 WEB" if models_cfg.shared_encoder else "asymmetric (WEB + SAT)" LOGGER.info("Building model — %s, %s", mode_str, enc_str) - if cfg.backbone == "sofia": - preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config} - if cfg.sofia_preset not in preset_map: - raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}") - sofia_cfg = preset_map[cfg.sofia_preset]() - sofia_cfg.d_descriptor = cfg.sofia_d_descriptor - sofia_cfg.text_film_dim = cfg.sofia_d_descriptor - sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode - sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode - sofia_cfg.mamba_variant = cfg.sofia_mamba_variant - sofia_cfg.mamba_backend = cfg.sofia_mamba_backend - sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge - sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations) + + if backbone == "sofia_v71": + assert isinstance(models_cfg, SOFIAv71ModelsConfig) + # Build SOFIAConfig from the gin-loaded SOFIAv71ModelsConfig. + # All architectural fields come from gin — no preset factory needed. + sofia_cfg = SOFIAConfig( + input_size=models_cfg.input_size, + in_channels=models_cfg.in_channels, + stem_mid=models_cfg.stem_mid, + stem_out=models_cfg.stem_out, + embed_dims=list(models_cfg.embed_dims), + depths=list(models_cfg.depths), + mbconv_expand=models_cfg.mbconv_expand, + se_ratio=models_cfg.se_ratio, + strip_kernel_s1=models_cfg.strip_kernel_s1, + strip_kernel_s2=models_cfg.strip_kernel_s2, + mix_kernels=list(models_cfg.mix_kernels), + use_dcn_strip=models_cfg.use_dcn_strip, + mamba_d_state=models_cfg.mamba_d_state, + mamba_dt_rank=models_cfg.mamba_dt_rank, + mamba_backend=models_cfg.mamba_backend, + mamba_variant=models_cfg.mamba_variant, + mamba_extra_kwargs=dict(models_cfg.mamba_extra_kwargs), + num_heads_s3=models_cfg.num_heads_s3, + num_heads_s4=models_cfg.num_heads_s4, + use_strip_branch_s3=models_cfg.use_strip_branch_s3, + use_strip_branch_s4=models_cfg.use_strip_branch_s4, + ffn_expand=models_cfg.ffn_expand, + use_evss_bridge=models_cfg.use_evss_bridge, + evss_bridge_locations=list(models_cfg.evss_bridge_locations), + neck_channels=models_cfg.neck_channels, + d_descriptor=models_cfg.d_descriptor, + use_asymmetric_heads=models_cfg.use_asymmetric_heads, + chp_rings=models_cfg.chp_rings, + chp_angles=models_cfg.chp_angles, + chp_harmonics=models_cfg.chp_harmonics, + use_film_altitude=models_cfg.use_film_altitude, + altitude_norm=models_cfg.altitude_norm, + ring_count=models_cfg.ring_count, + use_ring_aux=models_cfg.use_ring_aux, + return_normalized=models_cfg.return_normalized, + # Disable text fusion when baseline_mode is on, regardless of gin. + use_text_film_sat=models_cfg.use_text_film_sat and not models_common_cfg.baseline_mode, + use_text_film_uav=models_cfg.use_text_film_uav and not models_common_cfg.baseline_mode, + text_film_dim=models_cfg.text_film_dim, + text_film_hidden=models_cfg.text_film_hidden, + share_stages_1_2=models_cfg.share_stages_1_2, + enable_kd_taps=models_cfg.enable_kd_taps, + precision=models_cfg.precision, + ) model = SOFIAFusionEncoder( sofia_cfg=sofia_cfg, - lrsclip_path=cfg.lrsclip_path, - init_gate=cfg.init_gate, - baseline_mode=cfg.baseline_mode, - lora_rank=cfg.sofia_lora_rank, - device=cfg.device, - ).to(cfg.device) - elif cfg.backbone == "sofia_v1": + lrsclip_path=models_common_cfg.lrsclip_path, + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + lora_rank=models_cfg.lora_rank, + device=hardware_cfg.device, + ).to(hardware_cfg.device) + elif backbone == "sofia_v1": + assert isinstance(models_cfg, SOFIAv1ModelsConfig) sofia_v1_cfg = SOFIAv1Config( - variant=cfg.sofia_v1_variant, - dcn_variant=cfg.sofia_v1_dcn_variant, - d_descriptor=cfg.sofia_v1_d_descriptor, - text_film_dim=cfg.sofia_v1_d_descriptor, - use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode, - use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode, - use_film_altitude=cfg.sofia_v1_use_film_altitude, + variant=models_cfg.variant_label, + dcn_variant=models_cfg.dcn_variant, + d_descriptor=models_cfg.d_descriptor, + text_film_dim=models_cfg.d_descriptor, # match d_descriptor (preserves old behavior) + use_text_film_uav=models_cfg.use_text_film_uav and not models_common_cfg.baseline_mode, + use_text_film_sat=models_cfg.use_text_film_sat and not models_common_cfg.baseline_mode, + use_film_altitude=models_cfg.use_film_altitude, ) model = SOFIAv1FusionEncoder( sofia_cfg=sofia_v1_cfg, - lrsclip_path=cfg.lrsclip_path, - init_gate=cfg.init_gate, - baseline_mode=cfg.baseline_mode, - lora_rank=cfg.sofia_v1_lora_rank, - device=cfg.device, - ).to(cfg.device) - else: + lrsclip_path=models_common_cfg.lrsclip_path, + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + lora_rank=models_cfg.lora_rank, + device=hardware_cfg.device, + ).to(hardware_cfg.device) + elif backbone == "stripnet": + assert isinstance(models_cfg, StripNetModelsConfig) + # AsymmetricEncoder also handles StripNet — pass dummy DINO paths, + # they're not used when backbone='stripnet'. (DINO fields not + # bindable on StripNetModelsConfig — by design.) model = AsymmetricEncoder( - dino_web_path=cfg.dino_web_path, - dino_sat_path=cfg.dino_sat_path, - lrsclip_path=cfg.lrsclip_path, - init_gate=cfg.init_gate, - baseline_mode=cfg.baseline_mode, - shared_encoder=cfg.shared_encoder, - mona_bottleneck=cfg.mona_bottleneck, - mona_last_n_blocks=cfg.mona_last_n_blocks, - device=cfg.device, - backbone=cfg.backbone, - stripnet_path=cfg.stripnet_path, - stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages, - stripnet_freeze=cfg.stripnet_freeze, - ).to(cfg.device) + dino_web_path="nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth", + dino_sat_path="nn_models/DINO_SAT/model.safetensors", + lrsclip_path=models_common_cfg.lrsclip_path, + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + shared_encoder=True, # StripNet is always shared + mona_bottleneck=64, + mona_last_n_blocks=12, + device=hardware_cfg.device, + backbone=backbone, + stripnet_path=models_cfg.stripnet_path, + stripnet_mona_last_n_stages=models_cfg.stripnet_mona_last_n_stages, + stripnet_freeze=models_cfg.stripnet_freeze, + ).to(hardware_cfg.device) + else: # dinov3 + assert isinstance(models_cfg, DINOv3ModelsConfig) + model = AsymmetricEncoder( + dino_web_path=models_cfg.dino_web_path, + dino_sat_path=models_cfg.dino_sat_path, + lrsclip_path=models_common_cfg.lrsclip_path, + init_gate=models_common_cfg.init_gate, + baseline_mode=models_common_cfg.baseline_mode, + shared_encoder=models_cfg.shared_encoder, + mona_bottleneck=models_cfg.mona_bottleneck, + mona_last_n_blocks=models_cfg.mona_last_n_blocks, + device=hardware_cfg.device, + backbone=backbone, + stripnet_path="nn_models/STRIPNET/stripnet_s.pth", + stripnet_mona_last_n_stages=0, + stripnet_freeze=True, + ).to(hardware_cfg.device) LOGGER.info("embed_dim=%d", model.embed_dim) # --- Gradient checkpointing (trade compute for VRAM) --- # StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it. - if cfg.gradient_checkpointing and cfg.backbone == "dinov3": - if cfg.shared_encoder: + if hardware_cfg.gradient_checkpointing and backbone == "dinov3": + assert isinstance(models_cfg, DINOv3ModelsConfig) + if models_cfg.shared_encoder: model.image_encoder.set_gradient_checkpointing(True) else: model.drone_encoder.set_gradient_checkpointing(True) @@ -711,10 +739,10 @@ def train(cfg: TrainConfigGTAUAV) -> None: if model.text_encoder is not None: model.text_encoder.transformer.gradient_checkpointing = True LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)") - elif cfg.gradient_checkpointing and cfg.backbone in ("stripnet", "sofia", "sofia_v1"): + elif hardware_cfg.gradient_checkpointing and backbone in ("stripnet", "sofia_v71", "sofia_v1"): if model.text_encoder is not None: model.text_encoder.transformer.gradient_checkpointing = True - LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone) + LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", backbone) n_trainable = sum(p.numel() for p in model.trainable_parameters()) n_total = sum(p.numel() for p in model.parameters()) @@ -724,45 +752,56 @@ def train(cfg: TrainConfigGTAUAV) -> None: ) # --- Model summary (torchinfo) --- - model_summary = print_model_summary(model, device=cfg.device) + model_summary = print_model_summary(model, device=hardware_cfg.device) (output_dir / "model_summary.txt").write_text(model_summary) # --- W&B model watching (gradient + weight histograms) --- if tracker.has_wandb: tracker.watch_model(model, log_freq=50) - # Loss. - if cfg.loss_type == "symmetric": + # Loss. InfoNCELoss / WeightedInfoNCELoss are NOT @gin.configurable — + # all parameters arrive explicitly from training_cfg. + if training_cfg.loss_type == "symmetric": loss_fn = InfoNCELoss( - temperature_init=cfg.tau_init, - learnable_temperature=cfg.learnable_temperature, - label_smoothing=cfg.label_smoothing, - weight_q2g=cfg.weight_q2g, - weight_g2q=cfg.weight_g2q, + temperature_init=training_cfg.tau_init, + temperature_final=training_cfg.tau_final, + label_smoothing=training_cfg.label_smoothing, + weight_q2g=training_cfg.weight_q2g, + weight_g2q=training_cfg.weight_g2q, + learnable_temperature=training_cfg.learnable_temperature, + tau_min=training_cfg.tau_min, + tau_max=training_cfg.tau_max, + hard_mining_k=training_cfg.hard_mining_k, ) loss_name = "SymmetricInfoNCE" - elif cfg.loss_type == "weighted": + elif training_cfg.loss_type == "weighted": loss_fn = WeightedInfoNCELoss( - temperature_init=cfg.tau_init, - learnable_temperature=cfg.learnable_temperature, - label_smoothing=cfg.label_smoothing, + temperature_init=training_cfg.tau_init, + learnable_temperature=training_cfg.learnable_temperature, + label_smoothing=training_cfg.label_smoothing, + k=training_cfg.weighted_loss_k, + tau_min=training_cfg.tau_min, + tau_max=training_cfg.tau_max, ) loss_name = "WeightedInfoNCE" else: - raise ValueError(f"Unknown loss_type={cfg.loss_type!r} (expected 'symmetric' or 'weighted')") + raise ValueError( + f"Unknown loss_type={training_cfg.loss_type!r} " + f"(expected 'symmetric' or 'weighted')", + ) LOGGER.info( "Loss: %s Temperature: %s (init=%.3f) q2g=%.2f g2q=%.2f", loss_name, - "learnable" if cfg.learnable_temperature else "fixed", - cfg.tau_init, cfg.weight_q2g, cfg.weight_g2q, + "learnable" if training_cfg.learnable_temperature else "fixed", + training_cfg.tau_init, training_cfg.weight_q2g, training_cfg.weight_g2q, ) # Hard negative memory bank. neg_bank = None - if cfg.neg_bank_size > 0: - neg_bank = NegativeMemoryBank(size=cfg.neg_bank_size, dim=model.embed_dim).to(cfg.device) - LOGGER.info("Negative memory bank: size=%d, dim=%d", cfg.neg_bank_size, model.embed_dim) + if training_cfg.neg_bank_size > 0: + neg_bank = NegativeMemoryBank(size=training_cfg.neg_bank_size, dim=model.embed_dim).to(hardware_cfg.device) + LOGGER.info("Negative memory bank: size=%d, dim=%d", training_cfg.neg_bank_size, model.embed_dim) # Data — separate transforms for train (augmented) and eval (clean). drone_train_tf = get_drone_train_transform(image_size=256) @@ -770,42 +809,48 @@ def train(cfg: TrainConfigGTAUAV) -> None: eval_tf = get_dino_transform(image_size=256) train_ds = GTAUAVDataset( - pair_json=cfg.train_json, - rgb_root=cfg.rgb_root, - caption_root=cfg.caption_root, + pair_json=pipeline_cfg.train_json, + rgb_root=pipeline_cfg.rgb_root, + caption_root=pipeline_cfg.caption_root, drone_transform=drone_train_tf, sat_transform=sat_train_tf, - filter_meta=cfg.filter_meta, + filter_meta=pipeline_cfg.filter_meta, ) test_ds = GTAUAVDataset( - pair_json=cfg.test_json, - rgb_root=cfg.rgb_root, - caption_root=cfg.caption_root, + pair_json=pipeline_cfg.test_json, + rgb_root=pipeline_cfg.rgb_root, + caption_root=pipeline_cfg.caption_root, image_transform=eval_tf, - filter_meta=cfg.filter_meta, + filter_meta=pipeline_cfg.filter_meta, ) sat_cand_list = [entry["sat_candidates"] for entry in train_ds.entries] # Backward compat: `use_mutex_sampler=False` overrides to plain shuffle. - effective_sampler_type = cfg.sampler_type if cfg.use_mutex_sampler else "none" + effective_sampler_type = training_cfg.sampler_type if training_cfg.use_mutex_sampler else "none" if effective_sampler_type == "dss": batch_sampler = DynamicSimilaritySampler( - sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed, - knn_device=cfg.dss_knn_device, - use_lsh=cfg.dss_use_lsh, - lsh_num_tables=cfg.dss_lsh_num_tables, - lsh_num_bits=cfg.dss_lsh_num_bits, + sat_cand_list, + batch_size=hardware_cfg.batch_size, + shuffle=True, + seed=pipeline_cfg.seed, + knn_device=training_cfg.dss_knn_device, + use_lsh=training_cfg.dss_use_lsh, + lsh_num_tables=training_cfg.dss_lsh_num_tables, + lsh_num_bits=training_cfg.dss_lsh_num_bits, ) LOGGER.info( "Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs", - cfg.dss_knn_device, - " + LSH" if cfg.dss_use_lsh else "", - cfg.dss_warmup_epochs, cfg.dss_reembed_every, + training_cfg.dss_knn_device, + " + LSH" if training_cfg.dss_use_lsh else "", + training_cfg.dss_warmup_epochs, training_cfg.dss_reembed_every, ) elif effective_sampler_type == "mutex": batch_sampler = MutuallyExclusiveSampler( - sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed, + sat_cand_list, + batch_size=hardware_cfg.batch_size, + shuffle=True, + seed=pipeline_cfg.seed, ) LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch") else: @@ -816,76 +861,84 @@ def train(cfg: TrainConfigGTAUAV) -> None: train_loader = DataLoader( train_ds, batch_sampler=batch_sampler, - num_workers=cfg.num_workers, + num_workers=hardware_cfg.num_workers, collate_fn=collate_gtauav_batch, pin_memory=True, ) else: train_loader = DataLoader( train_ds, - batch_size=cfg.batch_size, + batch_size=hardware_cfg.batch_size, shuffle=True, - num_workers=cfg.num_workers, + num_workers=hardware_cfg.num_workers, collate_fn=collate_gtauav_batch, pin_memory=True, drop_last=True, ) emb_cache: EmbeddingCache | None = None - if cfg.dss_cache_dir is not None: - emb_cache = EmbeddingCache(cfg.dss_cache_dir) - LOGGER.info("DSS embedding cache: %s", cfg.dss_cache_dir) + if training_cfg.dss_cache_dir is not None: + emb_cache = EmbeddingCache(training_cfg.dss_cache_dir) + LOGGER.info("DSS embedding cache: %s", training_cfg.dss_cache_dir) test_loader = DataLoader( test_ds, - batch_size=cfg.batch_size, + batch_size=hardware_cfg.batch_size, shuffle=False, - num_workers=cfg.num_workers, + num_workers=hardware_cfg.num_workers, collate_fn=collate_gtauav_batch, pin_memory=True, ) # Train eval loader: clean transforms (no augmentation), for R@K on train set. train_eval_ds = GTAUAVDataset( - pair_json=cfg.train_json, - rgb_root=cfg.rgb_root, - caption_root=cfg.caption_root, + pair_json=pipeline_cfg.train_json, + rgb_root=pipeline_cfg.rgb_root, + caption_root=pipeline_cfg.caption_root, image_transform=eval_tf, - filter_meta=cfg.filter_meta, + filter_meta=pipeline_cfg.filter_meta, ) train_eval_loader = DataLoader( train_eval_ds, - batch_size=cfg.batch_size, + batch_size=hardware_cfg.batch_size, shuffle=False, - num_workers=cfg.num_workers, + num_workers=hardware_cfg.num_workers, collate_fn=collate_gtauav_batch, pin_memory=True, ) - effective_batch = cfg.batch_size * cfg.grad_accum_steps + effective_batch = hardware_cfg.batch_size * hardware_cfg.grad_accum_steps LOGGER.info( "train=%d test=%d batch=%d accum=%d effective_batch=%d", - len(train_ds), len(test_ds), cfg.batch_size, cfg.grad_accum_steps, effective_batch, + len(train_ds), len(test_ds), + hardware_cfg.batch_size, hardware_cfg.grad_accum_steps, effective_batch, ) - # Optimizer — per-group LR (text encoder gets lower LR). + # Optimizer — per-group LR (text encoder gets lower LR, StripNet backbone optionally). + stripnet_lr_factor = ( + models_cfg.stripnet_backbone_lr_factor + if isinstance(models_cfg, StripNetModelsConfig) + else 0.1 # default; not used unless StripNet group is non-empty + ) param_groups = _build_param_groups( - model, cfg.learning_rate, cfg.text_lr_factor, - stripnet_backbone_lr_factor=cfg.stripnet_backbone_lr_factor, + model, + training_cfg.learning_rate, + training_cfg.text_lr_factor, + stripnet_backbone_lr_factor=stripnet_lr_factor, ) # Include loss temperature if learnable. - if cfg.learnable_temperature and loss_fn.logit_scale is not None: + if training_cfg.learnable_temperature and loss_fn.logit_scale is not None: param_groups[0]["params"].append(loss_fn.logit_scale) - optimizer = AdamW(param_groups, weight_decay=cfg.weight_decay) + optimizer = AdamW(param_groups, weight_decay=training_cfg.weight_decay) - lr_info = f"proj={cfg.learning_rate:.0e}" - if not cfg.baseline_mode: - lr_info += f" text={cfg.learning_rate * cfg.text_lr_factor:.0e}" - LOGGER.info("Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, cfg.warmup_epochs) + lr_info = f"proj={training_cfg.learning_rate:.0e}" + if not models_common_cfg.baseline_mode: + lr_info += f" text={training_cfg.learning_rate * training_cfg.text_lr_factor:.0e}" + LOGGER.info("Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, pipeline_cfg.warmup_epochs) # Scheduler — cosine with linear warmup (counted in optimizer steps). - steps_per_epoch = math.ceil(len(train_loader) / cfg.grad_accum_steps) - total_steps = cfg.epochs * steps_per_epoch - warmup_steps = cfg.warmup_epochs * steps_per_epoch + steps_per_epoch = math.ceil(len(train_loader) / hardware_cfg.grad_accum_steps) + total_steps = pipeline_cfg.epochs * steps_per_epoch + warmup_steps = pipeline_cfg.warmup_epochs * steps_per_epoch with warnings.catch_warnings(): warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*") scheduler = LambdaLR( @@ -893,7 +946,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps), last_epoch=-1, ) - scaler = GradScaler(enabled=cfg.use_amp) + scaler = GradScaler(enabled=hardware_cfg.use_amp) # Restore optimizer/scheduler/loss state on resume. if resume_ckpt is not None: @@ -912,20 +965,20 @@ def train(cfg: TrainConfigGTAUAV) -> None: # --- Optional profiler (first epoch only) --- profiler = None - if cfg.use_profiler and start_epoch == 0: + if tracking_cfg.use_profiler and start_epoch == 0: profiler = TrainingProfiler( output_dir=output_dir, - n_warmup=cfg.profiler_warmup, - n_active=cfg.profiler_active, + n_warmup=tracking_cfg.profiler_warmup, + n_active=tracking_cfg.profiler_active, ) profiler.start() - LOGGER.info("Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch) + LOGGER.info("Starting training for %d epochs (from epoch %d)", pipeline_cfg.epochs, start_epoch) global_step = start_epoch * steps_per_epoch best_r1 = 0.0 - for epoch in range(start_epoch, cfg.epochs): + for epoch in range(start_epoch, pipeline_cfg.epochs): model.train() if batch_sampler is not None: batch_sampler.set_epoch(epoch) @@ -933,8 +986,8 @@ def train(cfg: TrainConfigGTAUAV) -> None: # DSS re-embedding: refresh query embeddings before the epoch starts. if ( isinstance(batch_sampler, DynamicSimilaritySampler) - and epoch >= cfg.dss_warmup_epochs - and (epoch - cfg.dss_warmup_epochs) % cfg.dss_reembed_every == 0 + and epoch >= training_cfg.dss_warmup_epochs + and (epoch - training_cfg.dss_warmup_epochs) % training_cfg.dss_reembed_every == 0 ): query_embs: torch.Tensor | None = None if emb_cache is not None: @@ -943,9 +996,9 @@ def train(cfg: TrainConfigGTAUAV) -> None: LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch) t_embed = time.time() query_embs = _embed_drone_queries( - model, train_ds, cfg.device, - batch_size=cfg.batch_size * cfg.grad_accum_steps, - num_workers=cfg.num_workers, + model, train_ds, hardware_cfg.device, + batch_size=hardware_cfg.batch_size * hardware_cfg.grad_accum_steps, + num_workers=hardware_cfg.num_workers, ) LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed) if emb_cache is not None: @@ -960,25 +1013,25 @@ def train(cfg: TrainConfigGTAUAV) -> None: pbar = tqdm( train_loader, - desc=f" Epoch {epoch + 1}/{cfg.epochs}", + desc=f" Epoch {epoch + 1}/{pipeline_cfg.epochs}", unit="batch", leave=False, ) - accum = cfg.grad_accum_steps + accum = hardware_cfg.grad_accum_steps for batch in pbar: # Zero gradients only at the start of each accumulation window. if n_batches % accum == 0: optimizer.zero_grad(set_to_none=True) - drone_img = batch["drone_img"].to(cfg.device, non_blocking=True) - sat_img = batch["sat_img"].to(cfg.device, non_blocking=True) + drone_img = batch["drone_img"].to(hardware_cfg.device, non_blocking=True) + sat_img = batch["sat_img"].to(hardware_cfg.device, non_blocking=True) altitude = batch.get("altitude") if altitude is not None: - altitude = altitude.to(cfg.device, non_blocking=True) + altitude = altitude.to(hardware_cfg.device, non_blocking=True) # Model forward in AMP (fp16 for DINOv3/DGTRS encoders). - with autocast(device_type="cuda", enabled=cfg.use_amp): - if cfg.baseline_mode: + with autocast(device_type="cuda", enabled=hardware_cfg.use_amp): + if models_common_cfg.baseline_mode: embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude) else: embeddings = model( @@ -997,12 +1050,12 @@ def train(cfg: TrainConfigGTAUAV) -> None: loss_kwargs = { "embeddings": embeddings, "epoch": epoch, - "total_epochs": cfg.epochs, + "total_epochs": pipeline_cfg.epochs, "queue_negatives": queue_neg, } if isinstance(loss_fn, WeightedInfoNCELoss): loss_kwargs["positive_weights"] = batch["positive_weights"].to( - cfg.device, non_blocking=True, + hardware_cfg.device, non_blocking=True, ) loss_dict = loss_fn(**loss_kwargs) @@ -1011,27 +1064,22 @@ def train(cfg: TrainConfigGTAUAV) -> None: total_loss = loss_dict["total"] / accum scaler.scale(total_loss).backward() - # Enqueue current gallery AFTER backward. The queue buffer is aliased - # into the autograd graph through `queue_neg` (a view returned by - # `NegativeMemoryBank.get_queue`), so modifying it before backward - # triggers "variable needed for gradient computation has been modified - # by an inplace operation". Enqueueing here is semantically identical - # — the next step's queue state is the same either way. + # Enqueue current gallery AFTER backward. if neg_bank is not None: neg_bank.enqueue(embeddings["gallery"].detach()) # Optimizer step only after accumulating `accum` micro-batches. is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader) if is_accum_step: - if cfg.grad_clip > 0: + if training_cfg.grad_clip > 0: scaler.unscale_(optimizer) nn.utils.clip_grad_norm_( model.trainable_parameters(), - max_norm=cfg.grad_clip, + max_norm=training_cfg.grad_clip, ) # --- Gradient monitoring (after unscale, before step) --- - if cfg.log_grad_norms and n_batches % (50 * accum) < accum: + if tracking_cfg.log_grad_norms and n_batches % (50 * accum) < accum: grad_norms = compute_gradient_norms(model, loss_fn) tracker.log_gradients(epoch, grad_norms, step=global_step) if n_batches < accum: @@ -1099,17 +1147,21 @@ def train(cfg: TrainConfigGTAUAV) -> None: # Evaluation. train_recall = {} - if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1: + if (epoch + 1) % pipeline_cfg.eval_every == 0 or epoch == pipeline_cfg.epochs - 1: # Train R@K (subset — same size as test set for speed). train_eval_batches = len(test_loader) train_recall = _evaluate( - model, train_eval_loader, cfg.device, - loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs, + model, train_eval_loader, hardware_cfg.device, + loss_fn=loss_fn, epoch=epoch, total_epochs=pipeline_cfg.epochs, max_batches=train_eval_batches, desc="eval-train", ) epoch_record["train_recall"] = train_recall csv_logger.log_train_recall(epoch, train_recall) - tracker.log_train(epoch, {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, step=global_step) + tracker.log_train( + epoch, + {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, + step=global_step, + ) # Log train metrics to CSV (includes recall/AP if eval ran this epoch). train_row = {**means} @@ -1145,8 +1197,8 @@ def train(cfg: TrainConfigGTAUAV) -> None: # Val R@K (full test set). val_metrics = _evaluate( - model, test_loader, cfg.device, - loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs, + model, test_loader, hardware_cfg.device, + loss_fn=loss_fn, epoch=epoch, total_epochs=pipeline_cfg.epochs, desc="eval-val", ) epoch_record["val"] = val_metrics @@ -1177,14 +1229,14 @@ def train(cfg: TrainConfigGTAUAV) -> None: ) # --- Grad-CAM visualization --- - if cfg.use_gradcam and (epoch + 1) % cfg.gradcam_every == 0: + if tracking_cfg.use_gradcam and (epoch + 1) % tracking_cfg.gradcam_every == 0: from src.training.gradcam import generate_gradcam_samples overlays = generate_gradcam_samples( model=model, dataloader=test_loader, - device=cfg.device, + device=hardware_cfg.device, output_dir=str(output_dir), - n_samples=cfg.gradcam_samples, + n_samples=tracking_cfg.gradcam_samples, epoch=epoch, ) # Log first few overlays to tracker. @@ -1207,15 +1259,16 @@ def train(cfg: TrainConfigGTAUAV) -> None: "model_state": model.state_dict(), "optimizer_state": optimizer.state_dict(), "loss_state": loss_fn.state_dict(), - "baseline_mode": cfg.baseline_mode, - "backbone": cfg.backbone, + "baseline_mode": models_common_cfg.baseline_mode, + "backbone": backbone, } - if cfg.backbone in ("sofia", "sofia_v1"): + if backbone in ("sofia_v71", "sofia_v1"): ckpt_obj["sofia_cfg"] = model.sofia_cfg - else: - ckpt_obj["shared_encoder"] = cfg.shared_encoder - ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck - ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks + elif isinstance(models_cfg, DINOv3ModelsConfig): + ckpt_obj["shared_encoder"] = models_cfg.shared_encoder + ckpt_obj["mona_bottleneck"] = models_cfg.mona_bottleneck + ckpt_obj["mona_last_n_blocks"] = models_cfg.mona_last_n_blocks + # StripNet has no extra arch flags worth saving here (params come from gin on resume). _atomic_save(obj=ckpt_obj, path=output_dir / f"ckpt_epoch{epoch:03d}.pt") LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch) @@ -1227,11 +1280,11 @@ def train(cfg: TrainConfigGTAUAV) -> None: # Save final eval report. LOGGER.info("Running final evaluation...") final_metrics = _evaluate( - model, test_loader, cfg.device, - loss_fn=loss_fn, epoch=cfg.epochs - 1, total_epochs=cfg.epochs, + model, test_loader, hardware_cfg.device, + loss_fn=loss_fn, epoch=pipeline_cfg.epochs - 1, total_epochs=pipeline_cfg.epochs, ) report = { - "config": vars(cfg), + "config": full_config, "metrics": final_metrics, "history": history, } @@ -1277,119 +1330,9 @@ def train(cfg: TrainConfigGTAUAV) -> None: ) -def main() -> None: - parser = argparse.ArgumentParser(description="GTA-UAV caption test training.") - parser.add_argument( - "--config", type=str, default=None, - help="Path to gin config file (e.g. conf/gtauav_balanced.gin).", - ) - parser.add_argument( - "--baseline", action="store_true", - help="Run baseline mode (no text).", - ) - parser.add_argument( - "--resume", type=str, default=None, - help="Path to checkpoint to resume training from.", - ) - parser.add_argument( - "--output-dir", type=str, default=None, - help="Override output directory.", - ) - parser.add_argument( - "--filter-meta", type=str, default=None, - help="Path to seg_filter.json for excluding bad images.", - ) - parser.add_argument( - "--batch-size", type=int, default=None, - help="Batch size.", - ) - parser.add_argument( - "--grad-accum", type=int, default=None, - help="Gradient accumulation steps (effective_batch = batch_size * accum).", - ) - parser.add_argument( - "--epochs", type=int, default=None, - help="Number of epochs.", - ) - parser.add_argument( - "--lr", type=float, default=None, - help="Learning rate for projections.", - ) - parser.add_argument( - "--text-lr-factor", type=float, default=None, - help="Text encoder LR = lr * factor (default 0.1 = 10x lower).", - ) - parser.add_argument( - "--warmup-epochs", type=int, default=None, - help="Linear warmup epochs.", - ) - parser.add_argument( - "--init-gate", type=float, default=None, - help="Initial gate value (image weight).", - ) - # Tracking flags. - parser.add_argument("--wandb", action="store_true", help="Enable W&B tracking.") - parser.add_argument("--no-tb", action="store_true", help="Disable TensorBoard.") - parser.add_argument("--gradcam", action="store_true", help="Enable Grad-CAM visualization.") - parser.add_argument("--profile", action="store_true", help="Enable PyTorch profiler (first epoch).") - parser.add_argument("--no-grad-norms", action="store_true", help="Disable gradient norm logging.") - # Gin overrides. - parser.add_argument( - "--gin-param", type=str, nargs="*", default=[], - help="Gin parameter overrides (e.g. 'TrainConfigGTAUAV.epochs=30').", - ) - args = parser.parse_args() - - # Parse gin config if provided. - if args.config is not None: - gin.parse_config_file(args.config) - if args.gin_param: - gin.parse_config(args.gin_param) - - # Create config (gin bindings apply via @gin.configurable). - cfg = TrainConfigGTAUAV() - - # CLI overrides take priority over gin. - if args.baseline: - cfg.baseline_mode = True - if args.resume is not None: - cfg.resume_from = args.resume - if args.batch_size is not None: - cfg.batch_size = args.batch_size - if args.grad_accum is not None: - cfg.grad_accum_steps = args.grad_accum - if args.epochs is not None: - cfg.epochs = args.epochs - if args.lr is not None: - cfg.learning_rate = args.lr - if args.text_lr_factor is not None: - cfg.text_lr_factor = args.text_lr_factor - if args.warmup_epochs is not None: - cfg.warmup_epochs = args.warmup_epochs - if args.init_gate is not None: - cfg.init_gate = args.init_gate - if args.filter_meta is not None: - cfg.filter_meta = args.filter_meta - - # Tracking overrides. - if args.wandb: - cfg.use_wandb = True - if args.no_tb: - cfg.use_tb = False - if args.gradcam: - cfg.use_gradcam = True - if args.profile: - cfg.use_profiler = True - if args.no_grad_norms: - cfg.log_grad_norms = False - - if args.output_dir is not None: - cfg.output_dir = args.output_dir - elif args.baseline and args.output_dir is None: - cfg.output_dir = "out/gtauav/baseline" - - train(cfg) - - +# Direct execution removed — entry point is src/main.py per REQUIREMENTS_GIN_STYLE.md §5. if __name__ == "__main__": - main() + raise SystemExit( + "Direct execution removed. Use: python -m src.main ", + ) + diff --git a/src/training/train_gtauav_old.py b/src/training/train_gtauav_old.py new file mode 100644 index 0000000..026a610 --- /dev/null +++ b/src/training/train_gtauav_old.py @@ -0,0 +1,1395 @@ +from __future__ import annotations + +"""Training loop for CVGL caption test on GTA-UAV-LR dataset. + +Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion. +Single InfoNCE loss: query(drone+text) vs gallery(satellite). + +Supports gin-config, W&B, TensorBoard, Grad-CAM, gradient monitoring, +PyTorch Profiler, and torchinfo model summary. +""" + +import argparse +import json +import logging +import math +import time +import warnings +from dataclasses import dataclass, field +from pathlib import Path + +import coloredlogs +import gin +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.amp import GradScaler, autocast +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader +from tqdm import tqdm + +from src.datasets.gtauav_dataset import ( + GTAUAVDataset, + GTAUAVDroneQuery, + GTAUAVSatGallery, + collate_drone_query, + collate_gtauav_batch, + collate_sat_gallery, +) +from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler +from src.datasets.embedding_cache import EmbeddingCache +from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler +from src.losses.multi_infonce import InfoNCELoss +from src.losses.weighted_infonce import WeightedInfoNCELoss +from src.losses.hard_negatives import NegativeMemoryBank +from src.training.plot_metrics import generate_plots +from src.training.trackers import ExperimentTracker +from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary +from src.training.profiling import TrainingProfiler, print_model_summary +from src.models.asymmetric_encoder import ( + AsymmetricEncoder, + get_dino_transform, + get_drone_train_transform, + get_satellite_train_transform, +) +from src.models.sofia_fusion_encoder import SOFIAFusionEncoder +from src.models.sofia_v1 import SOFIAv1Config +from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder +from src.models.sofia_v71 import ( + sofia_l_config, + sofia_m_config, + sofia_tiny_config, +) + +LOGGER = logging.getLogger("caption_test.train_gtauav") + +# Default paths. +_RGB_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR" +_CAPTION_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-captions" +_TRAIN_JSON = "meta/train_80.json" +_TEST_JSON = "meta/test_20.json" + +_DINO_WEB = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth" +_DINO_SAT = "nn_models/DINO_SAT/model.safetensors" +_LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt" + + +@gin.configurable(module="src.training.train_gtauav") +@dataclass +class TrainConfigGTAUAV: + """Training configuration for GTA-UAV experiment.""" + + # Data. + train_json: str = _TRAIN_JSON + test_json: str = _TEST_JSON + rgb_root: str = _RGB_ROOT + caption_root: str = _CAPTION_ROOT + filter_meta: str | None = None + + # Model. + dino_web_path: str = _DINO_WEB + dino_sat_path: str = _DINO_SAT + lrsclip_path: str = _LRSCLIP + init_gate: float = 0.7 + baseline_mode: bool = False + shared_encoder: bool = True # single DINOv3 WEB for both branches (simpler, half the params) + mona_bottleneck: int = 64 + mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks + gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch) + # StripNet backbone option (replaces DINOv3 when backbone="stripnet"). + backbone: str = "dinov3" # "dinov3", "stripnet", or "sofia" + stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth" + stripnet_mona_last_n_stages: int = 2 # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA) + stripnet_freeze: bool = True # If False, StripNet backbone is fully trainable (full fine-tune) + stripnet_backbone_lr_factor: float = 0.1 # Backbone LR = learning_rate * factor (only when unfrozen) + # SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint. + sofia_preset: str = "Tiny" # "Tiny" | "M" | "L" + sofia_d_descriptor: int = 1024 # retrieval space (1024 = match TextFusionMLP out_dim) + sofia_use_text_film_uav: bool = True # mid-level text-FiLM in UAV head + sofia_use_text_film_sat: bool = True # mid-level text-FiLM in SAT head + sofia_lora_rank: int = 4 + sofia_mamba_variant: str = "mamba2" # "mamba1" | "mamba2" | "efficient_vmamba" + sofia_mamba_backend: str = "auto" # "auto" | "torch" | "mamba_ssm" + # EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in). + sofia_use_evss_bridge: bool = False + sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"]) + # SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch. + sofia_v1_variant: str = "small" # "tiny_tiny" | "tiny" | "small" | "small_v2" + sofia_v1_dcn_variant: str = "v2" # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky) + sofia_v1_d_descriptor: int = 1024 + sofia_v1_use_text_film_uav: bool = True + sofia_v1_use_text_film_sat: bool = True + sofia_v1_use_film_altitude: bool = True + sofia_v1_lora_rank: int = 4 + + # Training. + resume_from: str | None = None # path to checkpoint for resuming + output_dir: str = "out/gtauav/with_text" + epochs: int = 10 + batch_size: int = 8 + num_workers: int = 4 + learning_rate: float = 1e-4 + text_lr_factor: float = 0.1 # text encoder LR = learning_rate * factor + weight_decay: float = 1e-4 + grad_clip: float = 1.0 + grad_accum_steps: int = 1 # gradient accumulation steps (effective_batch = batch_size * accum) + use_amp: bool = True + eval_every: int = 2 + warmup_epochs: int = 2 + seed: int = 42 + device: str = "cuda" + + # Loss. + loss_type: str = "symmetric" # "symmetric" (InfoNCE) or "weighted" (WeightedInfoNCE) + tau_init: float = 0.07 + label_smoothing: float = 0.1 + learnable_temperature: bool = True + weight_q2g: float = 0.6 + weight_g2q: float = 0.4 + neg_bank_size: int = 4096 # hard negative memory bank size (0 = disabled) + + # Sampling. + sampler_type: str = "mutex" # "mutex" (no false negatives) or "dss" (DSS + mutex) + dss_reembed_every: int = 1 # Re-embed train queries every N epochs for DSS. + dss_warmup_epochs: int = 1 # Use mutex-only for the first N epochs (fresh model embeddings aren't useful) + dss_knn_device: str = "cuda" # Device for similarity matmul in DSS sampler. + dss_use_lsh: bool = False # Approximate kNN via LSH (opt-in; exact is fast at 25K). + dss_lsh_num_tables: int = 8 + dss_lsh_num_bits: int = 14 + dss_cache_dir: str | None = None # Disk cache for embeddings; None = disabled. + # Legacy alias kept for backward compatibility. + use_mutex_sampler: bool = True + + # Tracking & diagnostics. + use_wandb: bool = False + use_tb: bool = True + wandb_project: str = "caption-test-gtauav" + wandb_run_name: str | None = None + wandb_entity: str | None = None + log_grad_norms: bool = True + use_gradcam: bool = False + gradcam_every: int = 5 # Grad-CAM every N epochs + gradcam_samples: int = 8 + use_profiler: bool = False + profiler_warmup: int = 3 + profiler_active: int = 5 + + +def _set_seed(seed: int) -> None: + import random as _random + import numpy as _np + _random.seed(seed) + _np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _atomic_save(obj: dict, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = path.with_suffix(path.suffix + ".tmp") + torch.save(obj, tmp_path) + tmp_path.replace(path) + + +def _build_param_groups( + model: nn.Module, + lr: float, + text_lr_factor: float, + stripnet_backbone_lr_factor: float = 0.1, +) -> list[dict]: + """Build optimizer param groups with separate LR for text encoder and unfrozen StripNet backbone. + + Groups: + - text_encoder.* → lr * text_lr_factor (default 1e-5) + - image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5) + - everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv, + SOFIA backbone+heads when backbone="sofia") → lr + """ + text_params = [] + backbone_params = [] + other_params = [] + + is_stripnet = isinstance(getattr(model, "image_encoder", None), nn.Module) and \ + getattr(model, "backbone", "dinov3") == "stripnet" + + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + if "text_encoder" in name: + text_params.append(param) + elif is_stripnet and name.startswith("image_encoder.backbone.") and "mona_" not in name: + backbone_params.append(param) + else: + other_params.append(param) + + groups = [{"params": other_params, "lr": lr}] + if text_params: + groups.append({"params": text_params, "lr": lr * text_lr_factor}) + if backbone_params: + groups.append({"params": backbone_params, "lr": lr * stripnet_backbone_lr_factor}) + + return groups + + +def _cosine_warmup_schedule( + warmup_steps: int, + total_steps: int, +) -> callable: + """Cosine annealing with linear warmup.""" + + def lr_lambda(step: int) -> float: + if step < warmup_steps: + return step / max(warmup_steps, 1) + progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + return lr_lambda + + +@torch.no_grad() +def _embed_drone_queries( + model: AsymmetricEncoder, + train_ds: GTAUAVDataset, + device: str, + batch_size: int, + num_workers: int, +) -> torch.Tensor: + """Forward all drone queries and return [N, D] embeddings on CPU. + + Used by DynamicSimilaritySampler to rank drones by visual similarity. + Runs with model.eval() but restores original train state afterwards. + """ + was_training = model.training + model.eval() + + query_ds = GTAUAVDroneQuery(train_ds) + loader = DataLoader( + query_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + collate_fn=collate_drone_query, + pin_memory=True, + ) + + embs: list[torch.Tensor] = [] + for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): + drone_img = batch["drone_img"].to(device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(device, non_blocking=True) + q = model.encode_query( + drone_img, + batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + altitude=altitude, + ) + embs.append(q.cpu()) + + if was_training: + model.train() + return torch.cat(embs, dim=0) + + +@torch.no_grad() +def _evaluate( + model: AsymmetricEncoder, + loader: DataLoader, + device: str, + loss_fn: nn.Module | None = None, + epoch: int = 0, + total_epochs: int = 1, + k_values: tuple[int, ...] = (1, 5, 10), + max_batches: int | None = None, + desc: str = "eval", +) -> dict[str, float]: + """Compute R@K and MRR on the full satellite gallery. + + Standard CVGL retrieval: forward every unique satellite in the dataset + once (gallery), forward every drone query, then rank gallery by + cosine similarity. A query counts as a hit@K if ANY of its valid + satellite matches (pair_pos_sate_img_list ∪ pair_pos_semipos_sate_img_list) + appears in the top-K. + + `max_batches` subsamples the drone queries (not the gallery) — useful + for a quick train-side sanity check. + """ + dataset = loader.dataset + if not isinstance(dataset, GTAUAVDataset): + raise TypeError(f"_evaluate expects GTAUAVDataset, got {type(dataset).__name__}") + + model.eval() + + batch_size = loader.batch_size or 32 + num_workers = getattr(loader, "num_workers", 0) + pin_memory = getattr(loader, "pin_memory", False) + + gallery_ds = GTAUAVSatGallery(dataset) + query_ds = GTAUAVDroneQuery(dataset) + + gallery_loader = DataLoader( + gallery_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=pin_memory, + collate_fn=collate_sat_gallery, + ) + query_loader = DataLoader( + query_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=pin_memory, + collate_fn=collate_drone_query, + ) + + # --- Gallery forward (all unique sats) --- + gallery_embs: list[torch.Tensor] = [] + gallery_names: list[str] = [] + for batch in tqdm(gallery_loader, desc=f" {desc}-gallery", unit="batch", leave=False): + sat_img = batch["sat_img"].to(device, non_blocking=True) + g = model.encode_gallery( + sat_img, + batch["sat_caption_l1"], batch["sat_caption_l2"], batch["sat_caption_l3"], + ) + gallery_embs.append(g.cpu()) + gallery_names.extend(batch["sat_names"]) + gallery = torch.cat(gallery_embs, dim=0) # [N_sat, D] + + # --- Query forward (optionally subsampled via max_batches) --- + query_embs: list[torch.Tensor] = [] + query_valid_names: list[list[str]] = [] + batch_losses: list[float] = [] + sat_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(gallery_names)} + + for i, batch in enumerate(tqdm(query_loader, desc=f" {desc}-query", unit="batch", leave=False)): + if max_batches is not None and i >= max_batches: + break + drone_img = batch["drone_img"].to(device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(device, non_blocking=True) + q = model.encode_query( + drone_img, + batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + altitude=altitude, + ) + query_embs.append(q.cpu()) + query_valid_names.extend(batch["valid_sat_names"]) + + # Per-batch loss: use first valid sat per query as its paired gallery. + if loss_fn is not None: + pair_indices: list[int] = [] + for names in batch["valid_sat_names"]: + for name in names: + if name in sat_name_to_idx: + pair_indices.append(sat_name_to_idx[name]) + break + else: + pair_indices.append(-1) + if all(idx >= 0 for idx in pair_indices): + paired_gallery = gallery[pair_indices].to(device) + fake_embeddings = { + "query": q, + "gallery": paired_gallery, + "gate_q": model.fusion_query.gate_value, + "gate_g": model.fusion_gallery.gate_value, + } + loss_dict = loss_fn(fake_embeddings, epoch=epoch, total_epochs=total_epochs) + batch_losses.append(float(loss_dict["total"].item())) + + query = torch.cat(query_embs, dim=0) # [N_q, D] + n_query = query.size(0) + + # --- Similarity + rankings --- + sim = query @ gallery.t() # [N_q, N_sat] + sorted_idx = sim.argsort(dim=1, descending=True) + + metrics: dict[str, float] = {} + if batch_losses: + metrics["loss"] = sum(batch_losses) / len(batch_losses) + + # Precompute valid gallery index sets per query. + valid_idx_per_query: list[set[int]] = [] + for names in query_valid_names: + valid = {sat_name_to_idx[n] for n in names if n in sat_name_to_idx} + valid_idx_per_query.append(valid) + + # R@K with multi-match. + for k in k_values: + hits = 0 + for i in range(n_query): + top_k = set(sorted_idx[i, :k].tolist()) + if valid_idx_per_query[i] & top_k: + hits += 1 + metrics[f"r@{k}_q2g"] = hits / max(n_query, 1) + + # MRR over valid matches (kept key `ap_q2g` for CSV/plot compatibility). + mrr_sum = 0.0 + n_scored = 0 + for i in range(n_query): + valid = valid_idx_per_query[i] + if not valid: + continue + n_scored += 1 + for rank, gidx in enumerate(sorted_idx[i].tolist()): + if gidx in valid: + mrr_sum += 1.0 / (rank + 1) + break + metrics["ap_q2g"] = mrr_sum / max(n_scored, 1) + + # --- g2q (satellite → drone): invert ground-truth --- + n_gallery = gallery.size(0) + valid_q_per_sat: list[set[int]] = [set() for _ in range(n_gallery)] + for q_idx, gset in enumerate(valid_idx_per_query): + for g_idx in gset: + valid_q_per_sat[g_idx].add(q_idx) + + sorted_idx_g2q = sim.t().argsort(dim=1, descending=True) # [N_sat, n_query] + n_scored_g2q = sum(1 for s in valid_q_per_sat if s) + + for k in k_values: + hits_g2q = 0 + for i in range(n_gallery): + valid = valid_q_per_sat[i] + if not valid: + continue + top_k = set(sorted_idx_g2q[i, :k].tolist()) + if valid & top_k: + hits_g2q += 1 + metrics[f"r@{k}_g2q"] = hits_g2q / max(n_scored_g2q, 1) + + mrr_sum_g2q = 0.0 + for i in range(n_gallery): + valid = valid_q_per_sat[i] + if not valid: + continue + for rank, qidx in enumerate(sorted_idx_g2q[i].tolist()): + if qidx in valid: + mrr_sum_g2q += 1.0 / (rank + 1) + break + metrics["ap_g2q"] = mrr_sum_g2q / max(n_scored_g2q, 1) + + metrics["n_query"] = float(n_query) + metrics["n_gallery"] = float(n_gallery) + metrics["n_scored_g2q"] = float(n_scored_g2q) + + metrics["gate_q"] = model.fusion_query.gate_value + metrics["gate_g"] = model.fusion_gallery.gate_value + return metrics + + +class CSVLogger: + """Log train/val metrics to CSV files using pandas. + + Creates: + {output_dir}/logs/train.csv — epoch-level train averages + {output_dir}/logs/val.csv — epoch-level val metrics + {output_dir}/logs/train_batches.csv — per-batch train metrics (all epochs) + {output_dir}/logs/epoch_{N}_batches.csv — per-batch for single epoch + """ + + def __init__(self, output_dir: Path) -> None: + self.log_dir = output_dir / "logs" + self.log_dir.mkdir(parents=True, exist_ok=True) + self._current_epoch: int = -1 + self._batch_columns: list[str] | None = None + self._cumulative_batch_path = self.log_dir / "train_batches.csv" + self._epoch_batch_path: Path | None = None + + # Load existing CSV data on resume (so plots show full history). + train_csv = self.log_dir / "train.csv" + val_csv = self.log_dir / "val.csv" + train_recall_csv = self.log_dir / "train_recall.csv" + if train_csv.exists(): + self.train_rows = pd.read_csv(train_csv).to_dict("records") + LOGGER.info("CSVLogger: loaded %d previous train epochs", len(self.train_rows)) + else: + self.train_rows = [] + if val_csv.exists(): + self.val_rows = pd.read_csv(val_csv).to_dict("records") + LOGGER.info("CSVLogger: loaded %d previous val epochs", len(self.val_rows)) + else: + self.val_rows = [] + if train_recall_csv.exists(): + self.train_recall_rows = pd.read_csv(train_recall_csv).to_dict("records") + else: + self.train_recall_rows = [] + + def log_batch(self, epoch: int, batch_idx: int, global_step: int, metrics: dict) -> None: + """Log metrics for a single training batch. Writes to disk immediately.""" + row = {"epoch": epoch, "batch": batch_idx, "global_step": global_step, **metrics} + + # On new epoch, start a fresh per-epoch CSV. + if epoch != self._current_epoch: + self._current_epoch = epoch + self._epoch_batch_path = self.log_dir / f"epoch_{epoch:03d}_batches.csv" + + # Determine columns on first call (consistent order). + if self._batch_columns is None: + self._batch_columns = list(row.keys()) + + row_df = pd.DataFrame([row], columns=self._batch_columns) + write_header = not self._cumulative_batch_path.exists() + + # Append to cumulative CSV. + row_df.to_csv( + self._cumulative_batch_path, mode="a", header=write_header, index=False, + ) + # Append to per-epoch CSV. + write_epoch_header = not self._epoch_batch_path.exists() + row_df.to_csv( + self._epoch_batch_path, mode="a", header=write_epoch_header, index=False, + ) + + def log_train(self, epoch: int, metrics: dict, lr: float, elapsed: float) -> None: + """Log epoch-level train averages. Replaces existing entry for same epoch on resume.""" + row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics} + # Remove previous entry for this epoch (resume may re-run it). + self.train_rows = [r for r in self.train_rows if r.get("epoch") != epoch] + self.train_rows.append(row) + pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False) + + def log_val(self, epoch: int, metrics: dict) -> None: + """Log val metrics. Replaces existing entry for same epoch on resume.""" + row = {"epoch": epoch, **metrics} + self.val_rows = [r for r in self.val_rows if r.get("epoch") != epoch] + self.val_rows.append(row) + pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False) + + def log_train_recall(self, epoch: int, metrics: dict) -> None: + """Log train recall metrics. Replaces existing entry for same epoch.""" + row = {"epoch": epoch, **metrics} + self.train_recall_rows = [r for r in self.train_recall_rows if r.get("epoch") != epoch] + self.train_recall_rows.append(row) + pd.DataFrame(self.train_recall_rows).to_csv(self.log_dir / "train_recall.csv", index=False) + + +def _clear_vram() -> None: + """Free VRAM from previous runs before starting.""" + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + allocated = torch.cuda.memory_allocated() / 1e9 + LOGGER.info("VRAM cleared. Current usage: %.2f GB", allocated) + + +def train(cfg: TrainConfigGTAUAV) -> None: + """Run full training loop.""" + coloredlogs.install( + level="INFO", + logger=LOGGER, + fmt="%(asctime)s %(name)s %(levelname)s %(message)s", + ) + _clear_vram() + _set_seed(cfg.seed) + output_dir = Path(cfg.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Save config. + with (output_dir / "config.json").open("w") as f: + json.dump(vars(cfg), f, indent=2) + + # --- Experiment tracker (W&B + TensorBoard) --- + tracker = ExperimentTracker( + output_dir=output_dir, + config=vars(cfg), + use_wandb=cfg.use_wandb, + use_tb=cfg.use_tb, + wandb_project=cfg.wandb_project, + wandb_run_name=cfg.wandb_run_name, + wandb_entity=cfg.wandb_entity, + ) + + # Model. + start_epoch = 0 + resume_ckpt = None + + if cfg.resume_from is not None: + LOGGER.info("Resuming from %s", cfg.resume_from) + if cfg.backbone == "sofia_v71": + model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint( + cfg.resume_from, + lrsclip_path=cfg.lrsclip_path, + device=cfg.device, + ) + elif cfg.backbone == "sofia_v1": + model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint( + cfg.resume_from, + lrsclip_path=cfg.lrsclip_path, + device=cfg.device, + ) + else: + model, resume_ckpt = AsymmetricEncoder.load_checkpoint( + cfg.resume_from, + dino_web_path=cfg.dino_web_path, + dino_sat_path=cfg.dino_sat_path, + lrsclip_path=cfg.lrsclip_path, + device=cfg.device, + ) + start_epoch = resume_ckpt.get("epoch", -1) + 1 + else: + mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)" + if cfg.backbone == "sofia_v71": + enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})" + elif cfg.backbone == "sofia_v1": + enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})" + elif cfg.backbone == "stripnet": + enc_str = "StripNet-small (shared, 512→1024 proj)" + else: + enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)" + LOGGER.info("Building model — %s, %s", mode_str, enc_str) + if cfg.backbone == "sofia_v71": + preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config} + if cfg.sofia_preset not in preset_map: + raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}") + sofia_cfg = preset_map[cfg.sofia_preset]() + sofia_cfg.d_descriptor = cfg.sofia_d_descriptor + sofia_cfg.text_film_dim = cfg.sofia_d_descriptor + sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode + sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode + sofia_cfg.mamba_variant = cfg.sofia_mamba_variant + sofia_cfg.mamba_backend = cfg.sofia_mamba_backend + sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge + sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations) + model = SOFIAFusionEncoder( + sofia_cfg=sofia_cfg, + lrsclip_path=cfg.lrsclip_path, + init_gate=cfg.init_gate, + baseline_mode=cfg.baseline_mode, + lora_rank=cfg.sofia_lora_rank, + device=cfg.device, + ).to(cfg.device) + elif cfg.backbone == "sofia_v1": + sofia_v1_cfg = SOFIAv1Config( + variant=cfg.sofia_v1_variant, + dcn_variant=cfg.sofia_v1_dcn_variant, + d_descriptor=cfg.sofia_v1_d_descriptor, + text_film_dim=cfg.sofia_v1_d_descriptor, + use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode, + use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode, + use_film_altitude=cfg.sofia_v1_use_film_altitude, + ) + model = SOFIAv1FusionEncoder( + sofia_cfg=sofia_v1_cfg, + lrsclip_path=cfg.lrsclip_path, + init_gate=cfg.init_gate, + baseline_mode=cfg.baseline_mode, + lora_rank=cfg.sofia_v1_lora_rank, + device=cfg.device, + ).to(cfg.device) + else: + model = AsymmetricEncoder( + dino_web_path=cfg.dino_web_path, + dino_sat_path=cfg.dino_sat_path, + lrsclip_path=cfg.lrsclip_path, + init_gate=cfg.init_gate, + baseline_mode=cfg.baseline_mode, + shared_encoder=cfg.shared_encoder, + mona_bottleneck=cfg.mona_bottleneck, + mona_last_n_blocks=cfg.mona_last_n_blocks, + device=cfg.device, + backbone=cfg.backbone, + stripnet_path=cfg.stripnet_path, + stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages, + stripnet_freeze=cfg.stripnet_freeze, + ).to(cfg.device) + LOGGER.info("embed_dim=%d", model.embed_dim) + + # --- Gradient checkpointing (trade compute for VRAM) --- + # StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it. + if cfg.gradient_checkpointing and cfg.backbone == "dinov3": + if cfg.shared_encoder: + model.image_encoder.set_gradient_checkpointing(True) + else: + model.drone_encoder.set_gradient_checkpointing(True) + model.sat_encoder.set_gradient_checkpointing(True) + if model.text_encoder is not None: + model.text_encoder.transformer.gradient_checkpointing = True + LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)") + elif cfg.gradient_checkpointing and cfg.backbone in ("stripnet", "sofia_v71", "sofia_v1"): + if model.text_encoder is not None: + model.text_encoder.transformer.gradient_checkpointing = True + LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone) + + n_trainable = sum(p.numel() for p in model.trainable_parameters()) + n_total = sum(p.numel() for p in model.parameters()) + LOGGER.info( + "trainable=%s (%.2f%%) total=%s", + f"{n_trainable:,}", 100.0 * n_trainable / max(n_total, 1), f"{n_total:,}", + ) + + # --- Model summary (torchinfo) --- + model_summary = print_model_summary(model, device=cfg.device) + (output_dir / "model_summary.txt").write_text(model_summary) + + # --- W&B model watching (gradient + weight histograms) --- + if tracker.has_wandb: + tracker.watch_model(model, log_freq=50) + + # Loss. + if cfg.loss_type == "symmetric": + loss_fn = InfoNCELoss( + temperature_init=cfg.tau_init, + learnable_temperature=cfg.learnable_temperature, + label_smoothing=cfg.label_smoothing, + weight_q2g=cfg.weight_q2g, + weight_g2q=cfg.weight_g2q, + ) + loss_name = "SymmetricInfoNCE" + elif cfg.loss_type == "weighted": + loss_fn = WeightedInfoNCELoss( + temperature_init=cfg.tau_init, + learnable_temperature=cfg.learnable_temperature, + label_smoothing=cfg.label_smoothing, + ) + loss_name = "WeightedInfoNCE" + else: + raise ValueError(f"Unknown loss_type={cfg.loss_type!r} (expected 'symmetric' or 'weighted')") + + LOGGER.info( + "Loss: %s Temperature: %s (init=%.3f) q2g=%.2f g2q=%.2f", + loss_name, + "learnable" if cfg.learnable_temperature else "fixed", + cfg.tau_init, cfg.weight_q2g, cfg.weight_g2q, + ) + + # Hard negative memory bank. + neg_bank = None + if cfg.neg_bank_size > 0: + neg_bank = NegativeMemoryBank(size=cfg.neg_bank_size, dim=model.embed_dim).to(cfg.device) + LOGGER.info("Negative memory bank: size=%d, dim=%d", cfg.neg_bank_size, model.embed_dim) + + # Data — separate transforms for train (augmented) and eval (clean). + drone_train_tf = get_drone_train_transform(image_size=256) + sat_train_tf = get_satellite_train_transform(image_size=256) + eval_tf = get_dino_transform(image_size=256) + + train_ds = GTAUAVDataset( + pair_json=cfg.train_json, + rgb_root=cfg.rgb_root, + caption_root=cfg.caption_root, + drone_transform=drone_train_tf, + sat_transform=sat_train_tf, + filter_meta=cfg.filter_meta, + ) + test_ds = GTAUAVDataset( + pair_json=cfg.test_json, + rgb_root=cfg.rgb_root, + caption_root=cfg.caption_root, + image_transform=eval_tf, + filter_meta=cfg.filter_meta, + ) + + sat_cand_list = [entry["sat_candidates"] for entry in train_ds.entries] + # Backward compat: `use_mutex_sampler=False` overrides to plain shuffle. + effective_sampler_type = cfg.sampler_type if cfg.use_mutex_sampler else "none" + + if effective_sampler_type == "dss": + batch_sampler = DynamicSimilaritySampler( + sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed, + knn_device=cfg.dss_knn_device, + use_lsh=cfg.dss_use_lsh, + lsh_num_tables=cfg.dss_lsh_num_tables, + lsh_num_bits=cfg.dss_lsh_num_bits, + ) + LOGGER.info( + "Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs", + cfg.dss_knn_device, + " + LSH" if cfg.dss_use_lsh else "", + cfg.dss_warmup_epochs, cfg.dss_reembed_every, + ) + elif effective_sampler_type == "mutex": + batch_sampler = MutuallyExclusiveSampler( + sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed, + ) + LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch") + else: + batch_sampler = None + LOGGER.info("Sampler: default shuffle (no mutex / no DSS)") + + if batch_sampler is not None: + train_loader = DataLoader( + train_ds, + batch_sampler=batch_sampler, + num_workers=cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + ) + else: + train_loader = DataLoader( + train_ds, + batch_size=cfg.batch_size, + shuffle=True, + num_workers=cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + drop_last=True, + ) + + emb_cache: EmbeddingCache | None = None + if cfg.dss_cache_dir is not None: + emb_cache = EmbeddingCache(cfg.dss_cache_dir) + LOGGER.info("DSS embedding cache: %s", cfg.dss_cache_dir) + test_loader = DataLoader( + test_ds, + batch_size=cfg.batch_size, + shuffle=False, + num_workers=cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + ) + # Train eval loader: clean transforms (no augmentation), for R@K on train set. + train_eval_ds = GTAUAVDataset( + pair_json=cfg.train_json, + rgb_root=cfg.rgb_root, + caption_root=cfg.caption_root, + image_transform=eval_tf, + filter_meta=cfg.filter_meta, + ) + train_eval_loader = DataLoader( + train_eval_ds, + batch_size=cfg.batch_size, + shuffle=False, + num_workers=cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + ) + + effective_batch = cfg.batch_size * cfg.grad_accum_steps + LOGGER.info( + "train=%d test=%d batch=%d accum=%d effective_batch=%d", + len(train_ds), len(test_ds), cfg.batch_size, cfg.grad_accum_steps, effective_batch, + ) + + # Optimizer — per-group LR (text encoder gets lower LR). + param_groups = _build_param_groups( + model, cfg.learning_rate, cfg.text_lr_factor, + stripnet_backbone_lr_factor=cfg.stripnet_backbone_lr_factor, + ) + # Include loss temperature if learnable. + if cfg.learnable_temperature and loss_fn.logit_scale is not None: + param_groups[0]["params"].append(loss_fn.logit_scale) + + optimizer = AdamW(param_groups, weight_decay=cfg.weight_decay) + + lr_info = f"proj={cfg.learning_rate:.0e}" + if not cfg.baseline_mode: + lr_info += f" text={cfg.learning_rate * cfg.text_lr_factor:.0e}" + LOGGER.info("Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, cfg.warmup_epochs) + + # Scheduler — cosine with linear warmup (counted in optimizer steps). + steps_per_epoch = math.ceil(len(train_loader) / cfg.grad_accum_steps) + total_steps = cfg.epochs * steps_per_epoch + warmup_steps = cfg.warmup_epochs * steps_per_epoch + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*") + scheduler = LambdaLR( + optimizer, + lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps), + last_epoch=-1, + ) + scaler = GradScaler(enabled=cfg.use_amp) + + # Restore optimizer/scheduler/loss state on resume. + if resume_ckpt is not None: + if "optimizer_state" in resume_ckpt: + optimizer.load_state_dict(resume_ckpt["optimizer_state"]) + LOGGER.info("Optimizer state restored") + if "loss_state" in resume_ckpt: + loss_fn.load_state_dict(resume_ckpt["loss_state"]) + LOGGER.info("Loss state restored (tau=%.4f)", loss_fn.current_temperature) + # Set scheduler last_epoch so it resumes at the correct LR. + scheduler.last_epoch = start_epoch * steps_per_epoch + LOGGER.info("Resuming from epoch %d", start_epoch) + + history: list[dict] = [] + csv_logger = CSVLogger(output_dir) + + # --- Optional profiler (first epoch only) --- + profiler = None + if cfg.use_profiler and start_epoch == 0: + profiler = TrainingProfiler( + output_dir=output_dir, + n_warmup=cfg.profiler_warmup, + n_active=cfg.profiler_active, + ) + profiler.start() + + LOGGER.info("Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch) + + global_step = start_epoch * steps_per_epoch + best_r1 = 0.0 + + for epoch in range(start_epoch, cfg.epochs): + model.train() + if batch_sampler is not None: + batch_sampler.set_epoch(epoch) + + # DSS re-embedding: refresh query embeddings before the epoch starts. + if ( + isinstance(batch_sampler, DynamicSimilaritySampler) + and epoch >= cfg.dss_warmup_epochs + and (epoch - cfg.dss_warmup_epochs) % cfg.dss_reembed_every == 0 + ): + query_embs: torch.Tensor | None = None + if emb_cache is not None: + query_embs = emb_cache.load(epoch) + if query_embs is None: + LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch) + t_embed = time.time() + query_embs = _embed_drone_queries( + model, train_ds, cfg.device, + batch_size=cfg.batch_size * cfg.grad_accum_steps, + num_workers=cfg.num_workers, + ) + LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed) + if emb_cache is not None: + emb_cache.save(epoch, query_embs) + t_sampler = time.time() + batch_sampler.update_embeddings(query_embs) + LOGGER.info("DSS: sampler update_embeddings took %.2fs", time.time() - t_sampler) + + epoch_start = time.time() + agg: dict[str, float] = {} + n_batches = 0 + + pbar = tqdm( + train_loader, + desc=f" Epoch {epoch + 1}/{cfg.epochs}", + unit="batch", + leave=False, + ) + accum = cfg.grad_accum_steps + for batch in pbar: + # Zero gradients only at the start of each accumulation window. + if n_batches % accum == 0: + optimizer.zero_grad(set_to_none=True) + + drone_img = batch["drone_img"].to(cfg.device, non_blocking=True) + sat_img = batch["sat_img"].to(cfg.device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(cfg.device, non_blocking=True) + + # Model forward in AMP (fp16 for DINOv3/DGTRS encoders). + with autocast(device_type="cuda", enabled=cfg.use_amp): + if cfg.baseline_mode: + embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude) + else: + embeddings = model( + drone_img=drone_img, + sat_img=sat_img, + caption_l1=batch["caption_l1"], + caption_l2=batch["caption_l2"], + caption_l3=batch["caption_l3"], + sat_caption_l1=batch["sat_caption_l1"], + sat_caption_l2=batch["sat_caption_l2"], + sat_caption_l3=batch["sat_caption_l3"], + altitude=altitude, + ) + # Loss — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights. + queue_neg = neg_bank.get_queue() if neg_bank is not None else None + loss_kwargs = { + "embeddings": embeddings, + "epoch": epoch, + "total_epochs": cfg.epochs, + "queue_negatives": queue_neg, + } + if isinstance(loss_fn, WeightedInfoNCELoss): + loss_kwargs["positive_weights"] = batch["positive_weights"].to( + cfg.device, non_blocking=True, + ) + loss_dict = loss_fn(**loss_kwargs) + + # Scale loss by accumulation steps so gradients average correctly. + raw_loss = float(loss_dict["total"].item()) # save before backward + total_loss = loss_dict["total"] / accum + scaler.scale(total_loss).backward() + + # Enqueue current gallery AFTER backward. The queue buffer is aliased + # into the autograd graph through `queue_neg` (a view returned by + # `NegativeMemoryBank.get_queue`), so modifying it before backward + # triggers "variable needed for gradient computation has been modified + # by an inplace operation". Enqueueing here is semantically identical + # — the next step's queue state is the same either way. + if neg_bank is not None: + neg_bank.enqueue(embeddings["gallery"].detach()) + + # Optimizer step only after accumulating `accum` micro-batches. + is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader) + if is_accum_step: + if cfg.grad_clip > 0: + scaler.unscale_(optimizer) + nn.utils.clip_grad_norm_( + model.trainable_parameters(), + max_norm=cfg.grad_clip, + ) + + # --- Gradient monitoring (after unscale, before step) --- + if cfg.log_grad_norms and n_batches % (50 * accum) < accum: + grad_norms = compute_gradient_norms(model, loss_fn) + tracker.log_gradients(epoch, grad_norms, step=global_step) + if n_batches < accum: + log_gradient_summary(grad_norms) + + scaler.step(optimizer) + scaler.update() + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*") + scheduler.step() + global_step += 1 + + # --- Per-batch tracking (log unscaled loss) --- + step_metrics = { + "loss": raw_loss, + "temperature": float(loss_dict["temperature"].item()), + "gate_q": float(loss_dict["gate_q"].item()), + "gate_g": float(loss_dict["gate_g"].item()), + "lr": optimizer.param_groups[0]["lr"], + } + tracker.log_train(epoch, step_metrics, step=global_step) + csv_logger.log_batch(epoch, n_batches, global_step, step_metrics) + + for key, val in loss_dict.items(): + agg[key] = agg.get(key, 0.0) + float(val.item()) + n_batches += 1 + + pbar.set_postfix( + loss=f"{raw_loss:.3f}", + tau=f"{step_metrics['temperature']:.4f}", + gq=f"{step_metrics['gate_q']:.3f}", + gg=f"{step_metrics['gate_g']:.3f}", + ) + + # --- Profiler step --- + if profiler is not None: + profiler.step() + if profiler.is_done(n_batches): + profiler.export() + profiler = None + + elapsed = time.time() - epoch_start + + means = {k: v / max(n_batches, 1) for k, v in agg.items()} + LOGGER.info( + "epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f", + epoch, elapsed, + optimizer.param_groups[0]["lr"], + means.get("total", 0.0), + means.get("temperature", 0.0), + means.get("gate_q", 1.0), + means.get("gate_g", 1.0), + ) + + epoch_record: dict = { + "epoch": epoch, + "elapsed_seconds": elapsed, + "train": means, + } + + # --- Log VRAM usage --- + if torch.cuda.is_available(): + vram_gb = torch.cuda.max_memory_allocated() / 1e9 + tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step) + + # Evaluation. + train_recall = {} + if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1: + # Train R@K (subset — same size as test set for speed). + train_eval_batches = len(test_loader) + train_recall = _evaluate( + model, train_eval_loader, cfg.device, + loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs, + max_batches=train_eval_batches, desc="eval-train", + ) + epoch_record["train_recall"] = train_recall + csv_logger.log_train_recall(epoch, train_recall) + tracker.log_train(epoch, {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, step=global_step) + + # Log train metrics to CSV (includes recall/AP if eval ran this epoch). + train_row = {**means} + if "total" in train_row: + train_row["train_loss"] = train_row.pop("total") + if train_recall: + train_row["r@1_q2g"] = train_recall.get("r@1_q2g", 0.0) + train_row["r@5_q2g"] = train_recall.get("r@5_q2g", 0.0) + train_row["r@10_q2g"] = train_recall.get("r@10_q2g", 0.0) + train_row["ap_q2g"] = train_recall.get("ap_q2g", 0.0) + train_row["r@1_g2q"] = train_recall.get("r@1_g2q", 0.0) + train_row["r@5_g2q"] = train_recall.get("r@5_g2q", 0.0) + train_row["r@10_g2q"] = train_recall.get("r@10_g2q", 0.0) + train_row["ap_g2q"] = train_recall.get("ap_g2q", 0.0) + csv_logger.log_train(epoch, train_row, optimizer.param_groups[0]["lr"], elapsed) + generate_plots(csv_logger.log_dir) + + if train_recall: + LOGGER.info( + "train-recall epoch=%d q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f " + "g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f", + epoch, + train_recall.get("r@1_q2g", 0.0), + train_recall.get("r@5_q2g", 0.0), + train_recall.get("r@10_q2g", 0.0), + train_recall.get("ap_q2g", 0.0), + train_recall.get("r@1_g2q", 0.0), + train_recall.get("r@5_g2q", 0.0), + train_recall.get("r@10_g2q", 0.0), + train_recall.get("ap_g2q", 0.0), + train_recall.get("loss", 0.0), + ) + + # Val R@K (full test set). + val_metrics = _evaluate( + model, test_loader, cfg.device, + loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs, + desc="eval-val", + ) + epoch_record["val"] = val_metrics + csv_logger.log_val(epoch, val_metrics) + generate_plots(csv_logger.log_dir) + tracker.log_val(epoch, val_metrics, step=global_step) + + # Track best R@1. + r1 = val_metrics.get("r@1_q2g", 0.0) + if r1 > best_r1: + best_r1 = r1 + tracker.log_scalar("val/best_r@1_q2g", best_r1, step=global_step) + + LOGGER.info( + "val epoch=%d q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f " + "g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f gate_q=%.4f", + epoch, + val_metrics.get("r@1_q2g", 0.0), + val_metrics.get("r@5_q2g", 0.0), + val_metrics.get("r@10_q2g", 0.0), + val_metrics.get("ap_q2g", 0.0), + val_metrics.get("r@1_g2q", 0.0), + val_metrics.get("r@5_g2q", 0.0), + val_metrics.get("r@10_g2q", 0.0), + val_metrics.get("ap_g2q", 0.0), + val_metrics.get("loss", 0.0), + val_metrics.get("gate_q", 1.0), + ) + + # --- Grad-CAM visualization --- + if cfg.use_gradcam and (epoch + 1) % cfg.gradcam_every == 0: + from src.training.gradcam import generate_gradcam_samples + overlays = generate_gradcam_samples( + model=model, + dataloader=test_loader, + device=cfg.device, + output_dir=str(output_dir), + n_samples=cfg.gradcam_samples, + epoch=epoch, + ) + # Log first few overlays to tracker. + for i, overlay in enumerate(overlays[:4]): + kind = "drone" if i % 2 == 0 else "sat" + tracker.log_image( + f"gradcam/{kind}_{i//2}", + overlay, + step=global_step, + caption=f"Epoch {epoch} {kind} Grad-CAM", + ) + + history.append(epoch_record) + + # Save checkpoint. Model architecture flags go into the ckpt so + # `AsymmetricEncoder.load_checkpoint` (or `SOFIAFusionEncoder.load_checkpoint`) + # can rebuild the right shape. + ckpt_obj = { + "epoch": epoch, + "model_state": model.state_dict(), + "optimizer_state": optimizer.state_dict(), + "loss_state": loss_fn.state_dict(), + "baseline_mode": cfg.baseline_mode, + "backbone": cfg.backbone, + } + if cfg.backbone in ("sofia_v71", "sofia_v1"): + ckpt_obj["sofia_cfg"] = model.sofia_cfg + else: + ckpt_obj["shared_encoder"] = cfg.shared_encoder + ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck + ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks + _atomic_save(obj=ckpt_obj, path=output_dir / f"ckpt_epoch{epoch:03d}.pt") + LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch) + + # Save history. + history_path = output_dir / "history.json" + with history_path.open("w", encoding="utf-8") as f: + json.dump(history, f, indent=2) + + # Save final eval report. + LOGGER.info("Running final evaluation...") + final_metrics = _evaluate( + model, test_loader, cfg.device, + loss_fn=loss_fn, epoch=cfg.epochs - 1, total_epochs=cfg.epochs, + ) + report = { + "config": vars(cfg), + "metrics": final_metrics, + "history": history, + } + report_path = output_dir / "eval_report.json" + with report_path.open("w", encoding="utf-8") as f: + json.dump(report, f, indent=2) + + # --- Log final summary to W&B --- + tracker.log_summary({ + "best_r@1_q2g": best_r1, + "final_r@1_q2g": final_metrics.get("r@1_q2g", 0.0), + "final_r@5_q2g": final_metrics.get("r@5_q2g", 0.0), + "final_r@10_q2g": final_metrics.get("r@10_q2g", 0.0), + "final_ap_q2g": final_metrics.get("ap_q2g", 0.0), + "final_r@1_g2q": final_metrics.get("r@1_g2q", 0.0), + "final_r@5_g2q": final_metrics.get("r@5_g2q", 0.0), + "final_r@10_g2q": final_metrics.get("r@10_g2q", 0.0), + "final_ap_g2q": final_metrics.get("ap_g2q", 0.0), + "final_gate_q": final_metrics.get("gate_q", 1.0), + "final_gate_g": final_metrics.get("gate_g", 1.0), + }) + + # --- Cleanup profiler if still running --- + if profiler is not None: + profiler.export() + + tracker.close() + + LOGGER.info("Training complete. Report: %s", report_path) + LOGGER.info( + "Final — q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f " + "g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f gate_q=%.4f gate_g=%.4f", + final_metrics.get("r@1_q2g", 0.0), + final_metrics.get("r@5_q2g", 0.0), + final_metrics.get("r@10_q2g", 0.0), + final_metrics.get("ap_q2g", 0.0), + final_metrics.get("r@1_g2q", 0.0), + final_metrics.get("r@5_g2q", 0.0), + final_metrics.get("r@10_g2q", 0.0), + final_metrics.get("ap_g2q", 0.0), + final_metrics.get("gate_q", 1.0), + final_metrics.get("gate_g", 1.0), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="GTA-UAV caption test training.") + parser.add_argument( + "--config", type=str, default=None, + help="Path to gin config file (e.g. conf/gtauav_balanced.gin).", + ) + parser.add_argument( + "--baseline", action="store_true", + help="Run baseline mode (no text).", + ) + parser.add_argument( + "--resume", type=str, default=None, + help="Path to checkpoint to resume training from.", + ) + parser.add_argument( + "--output-dir", type=str, default=None, + help="Override output directory.", + ) + parser.add_argument( + "--filter-meta", type=str, default=None, + help="Path to seg_filter.json for excluding bad images.", + ) + parser.add_argument( + "--batch-size", type=int, default=None, + help="Batch size.", + ) + parser.add_argument( + "--grad-accum", type=int, default=None, + help="Gradient accumulation steps (effective_batch = batch_size * accum).", + ) + parser.add_argument( + "--epochs", type=int, default=None, + help="Number of epochs.", + ) + parser.add_argument( + "--lr", type=float, default=None, + help="Learning rate for projections.", + ) + parser.add_argument( + "--text-lr-factor", type=float, default=None, + help="Text encoder LR = lr * factor (default 0.1 = 10x lower).", + ) + parser.add_argument( + "--warmup-epochs", type=int, default=None, + help="Linear warmup epochs.", + ) + parser.add_argument( + "--init-gate", type=float, default=None, + help="Initial gate value (image weight).", + ) + # Tracking flags. + parser.add_argument("--wandb", action="store_true", help="Enable W&B tracking.") + parser.add_argument("--no-tb", action="store_true", help="Disable TensorBoard.") + parser.add_argument("--gradcam", action="store_true", help="Enable Grad-CAM visualization.") + parser.add_argument("--profile", action="store_true", help="Enable PyTorch profiler (first epoch).") + parser.add_argument("--no-grad-norms", action="store_true", help="Disable gradient norm logging.") + # Gin overrides. + parser.add_argument( + "--gin-param", type=str, nargs="*", default=[], + help="Gin parameter overrides (e.g. 'TrainConfigGTAUAV.epochs=30').", + ) + args = parser.parse_args() + + # Parse gin config if provided. + if args.config is not None: + gin.parse_config_file(args.config) + if args.gin_param: + gin.parse_config(args.gin_param) + + # Create config (gin bindings apply via @gin.configurable). + cfg = TrainConfigGTAUAV() + + # CLI overrides take priority over gin. + if args.baseline: + cfg.baseline_mode = True + if args.resume is not None: + cfg.resume_from = args.resume + if args.batch_size is not None: + cfg.batch_size = args.batch_size + if args.grad_accum is not None: + cfg.grad_accum_steps = args.grad_accum + if args.epochs is not None: + cfg.epochs = args.epochs + if args.lr is not None: + cfg.learning_rate = args.lr + if args.text_lr_factor is not None: + cfg.text_lr_factor = args.text_lr_factor + if args.warmup_epochs is not None: + cfg.warmup_epochs = args.warmup_epochs + if args.init_gate is not None: + cfg.init_gate = args.init_gate + if args.filter_meta is not None: + cfg.filter_meta = args.filter_meta + + # Tracking overrides. + if args.wandb: + cfg.use_wandb = True + if args.no_tb: + cfg.use_tb = False + if args.gradcam: + cfg.use_gradcam = True + if args.profile: + cfg.use_profiler = True + if args.no_grad_norms: + cfg.log_grad_norms = False + + if args.output_dir is not None: + cfg.output_dir = args.output_dir + elif args.baseline and args.output_dir is None: + cfg.output_dir = "out/gtauav/baseline" + + train(cfg) + + +if __name__ == "__main__": + main() diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..c6d59bc --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1,2 @@ +"""Utilities: paths, seeding, IO.""" + diff --git a/src/utils/io_utils.py b/src/utils/io_utils.py new file mode 100644 index 0000000..e86276d --- /dev/null +++ b/src/utils/io_utils.py @@ -0,0 +1,51 @@ +"""IO helpers: atomic checkpoint saves, VRAM cleanup.""" + +from __future__ import annotations + +import gc +import logging +import os +import tempfile +from pathlib import Path +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + + +def atomic_save_torch(obj: Any, path: Path) -> None: + """Save a PyTorch object atomically via temp file + os.replace. + + On any failure (KeyboardInterrupt / SIGTERM included), the temp file + is removed. Makes --resume safe: a partial checkpoint never lands at + the destination path. + + Args: + obj: Anything torch.save can handle. + path: Destination path. Parent directory is created if missing. + """ + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(suffix=".pt.tmp", dir=path.parent) + os.close(fd) + try: + torch.save(obj, tmp) + os.replace(tmp, path) + except BaseException: + if os.path.exists(tmp): + os.remove(tmp) + raise + + +def clear_vram() -> None: + """Free VRAM and reset peak memory stats.""" + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + allocated_gb = torch.cuda.memory_allocated() / 1e9 + logger.info("VRAM cleared. Current usage: %.2f GB", allocated_gb) + + + + \ No newline at end of file diff --git a/src/utils/path_utils.py b/src/utils/path_utils.py new file mode 100644 index 0000000..ea0bebd --- /dev/null +++ b/src/utils/path_utils.py @@ -0,0 +1,32 @@ +"""Project root resolution via marker files.""" + +from __future__ import annotations + +from pathlib import Path + +# Markers identifying the project root (per REQUIREMENTS_GIN_STYLE.md §5). +_MARKERS: tuple[str, ...] = ("pyproject.toml", ".git", "in") + + +def get_proj_dir() -> str: + """Return absolute project root with trailing slash. + + Walks up from this file's directory until finding pyproject.toml, + .git, or in/. Searches up to 10 levels. + + Returns: + Project root path with trailing slash, e.g. '/home/user/caption-test/'. + + Raises: + RuntimeError: If no marker found within 10 parent directories. + """ + current = Path(__file__).resolve().parent + for _ in range(10): + if any((current / m).exists() for m in _MARKERS): + return str(current) + "/" + current = current.parent + raise RuntimeError( + f"Project root not found. Looked for {_MARKERS} starting at " + f"{Path(__file__).resolve().parent}", + ) + diff --git a/src/utils/seed_utils.py b/src/utils/seed_utils.py new file mode 100644 index 0000000..afe50cd --- /dev/null +++ b/src/utils/seed_utils.py @@ -0,0 +1,23 @@ +"""RNG seeding for reproducibility.""" + +from __future__ import annotations + +import random + +import numpy as np +import torch + + +def set_seed(seed: int = 42) -> None: + """Fix all RNG seeds (Python random, NumPy, PyTorch CPU + all CUDA devices). + + Args: + seed: Integer seed. + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + + \ No newline at end of file