checkpoint 15_05

CLAUDE.md

@@ -175,20 +175,20 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization.
 
 ## Датасет: GTA-UAV-LR
 
-- **RGB:** `/home/servml/Документы/datasets/GTA-UAV-LR/`
+- **RGB:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR/`
   - Drone: 33,763 PNG (512x384), altitudes 100-600m
   - Satellite: 14,640 PNG (256x256 RGBA)
   - Pairs: `cross-area-drone2sate-{train,test}.json` (primary split)
   - Metadata: `*_drone_meta.csv` (height, yaw, roll, pitch)
   - Origin: GTA V simulation (Los Santos)
 
-- **Captions:** `/home/servml/Документы/datasets/GTA-UAV-LR-captions/`
+- **Captions:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions/`
   - Drone: 33,411 JSON (32,635 multi-paragraph P1/P2/P3 + 776 short water-only)
   - Satellite: 6,546 JSON (все multi-paragraph)
   - Формат: 3 абзаца (P1 Inventory + P2 Spatial + P3 Fingerprint)
   - Token counts: ~430 output tokens per caption
 
-- **Segmentation:** `/home/servml/Документы/datasets/GTA-UAV-LR-aug/`
+- **Segmentation:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/`
   - 48,403 images, 17 классов (background, building, road, vegetation, water, ...)
   - Modalities: segm/, depth/, edge/, chm/, safetensors/
   - Query: 512x512, DB: 256x256
@@ -331,7 +331,7 @@ python -m scripts.compare_runs \
 
 ### UAV-VisLoc Prepare
 - **Путь:** `/home/servml/Документы/code/Yaroslav/UAV-VisLoc-prepare/scripts/prepare_dataset.py`
-- **Статус:** выполнен (2026-04-17), данные в `/home/servml/Документы/datasets/UAV_VisLoc_processed/` (25 GB)
+- **Статус:** выполнен (2026-04-17), данные в `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_processed/` (25 GB)
 - **Задача:** нарезка satellite кропов 512x512, stride 256 + resize drone -> 512x512
 - **Подробности:** см. ниже
 
@@ -376,7 +376,7 @@ Binary masks — natural FiLM gates. Modality dropout: text 0.3, CHM 0.5, seg 0.
 ## Датасеты (справочник)
 
 ### UAV-VisLoc
-- **Путь:** `/home/servml/Документы/datasets/UAV_VisLoc_dataset/`
+- **Путь:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_dataset/`
 - **Структура:** 11 маршрутов (папки `01`-`11`), каждая содержит:
   - `drone/` — drone-снимки (`XX_NNNN.JPG`)
   - `satelliteXX.tif` — спутниковая карта
@@ -403,8 +403,8 @@ Binary masks — natural FiLM gates. Modality dropout: text 0.3, CHM 0.5, seg 0.
 ### Запуск
 ```bash
 python scripts/prepare_dataset.py \
-    --src /home/servml/Документы/datasets/UAV_VisLoc_dataset \
-    --dst /home/servml/Документы/datasets/UAV_VisLoc_processed \
+    --src /media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_dataset \
+    --dst /media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_processed \
     --crop-size 512 --stride 256 --target-size 512
 ```
 

GTA-UAV-LR_Dataset_Analysis.md

@@ -2,8 +2,8 @@
 
 **Дата анализа:** 2026-04-21
 **Метод:** Эмпирический анализ данных на диске + статья arXiv:2409.16925 + GitHub-репозиторий авторов
-**Путь к данным:** `/home/servml/Документы/datasets/GTA-UAV-LR/`
-**Путь к аугментациям:** `/home/servml/Документы/datasets/GTA-UAV-LR-aug/`
+**Путь к данным:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR/`
+**Путь к аугментациям:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/`
 
 ---
 
@@ -284,7 +284,7 @@
 
 ## 8. АУГМЕНТИРОВАННЫЙ НАБОР (GTA-UAV-LR-aug)
 
-**Путь:** `/home/servml/Документы/datasets/GTA-UAV-LR-aug/`
+**Путь:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/`
 **Объём:** ~71 GB
 
 ### 8.1. Сгенерированные модальности

README.md

@@ -377,9 +377,9 @@ that list it as a valid candidate. `n_scored_g2q` is reported in metrics for tra
 ### V3 — GTA-UAV + DINOv3 + DGTRS-CLIP (active)
 
 **Dataset:** GTA-UAV-LR (33K drone + 14K satellite, GTA V synthetic)
-- RGB: `/home/servml/Документы/datasets/GTA-UAV-LR/`
-- Captions: `/home/servml/Документы/datasets/GTA-UAV-LR-captions/` (40K JSON, 3-paragraph VLM)
-- Segmentation: `/home/servml/Документы/datasets/GTA-UAV-LR-aug/` (17 classes)
+- RGB: `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR/`
+- Captions: `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions/` (40K JSON, 3-paragraph VLM)
+- Segmentation: `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/` (17 classes)
 - Seg filter: 37,498 passed / 10,905 excluded (>=90% background+water)
 - Split: 80/20 random (26,966 train / 6,742 test → 24,891/6,252 after seg filter)
 
@@ -584,3 +584,21 @@ tensorboard --logdir out/gtauav/with_text/tb_logs
 - Type hints on all signatures
 - Google-style docstrings
 - English-only comments
+
+# Authors 
+
+## Лицензия и автор
+
+**Автор:** Пикалёв Ярослав Сергеевич, к.т.н., ЛИСАД, ФГБНУ «ИПИИ»
+**Email:** i@pikaliov.ru
+
+**Рефакторинг ветка:** Павленко Богдан Викторович, м.н.с., ЛИСАД, ФГБНУ «ИПИИ»
+**Email:** bogdanpavl2000@mail.ru
+
+Проект реализован в рамках: (НИР FREN-2024-0002 «Извлечение семантической информации из изображений для автономных систем навигации БПЛА».)
+
+==**Код предоставляется для научных целей в рамках выполняемой работы.**==
+
+**Работы, ссылка на которые обязательна при использовании этой реализации:**
+1. Пикалёв Я. С., Павленко Б. В. Регрессионная нейронная сеть на основе регуляризации функции потерь… // Информатика и автоматизация (SPCRAS). 2025.
+2. Павленко Б. В., Пикалёв Я. С. Методика создания набора аэрофотоснимков для CVGL // ПИИ. 2024. № 4 (35). С. 101–112.

UAV-VisLoc_Dataset_Analysis.md

@@ -2,7 +2,7 @@
 
 **Дата анализа:** 2026-04-17
 **Метод:** Эмпирический анализ данных на диске + статья arXiv:2405.11936 + GitHub-репозиторий авторов
-**Путь к данным:** `/home/servml/Документы/datasets/UAV_VisLoc_dataset/`
+**Путь к данным:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_dataset/`
 
 ---
 

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/<name>/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 <preset_name>",
+    )
+```
+
+---
+
+## 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/<name>/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/<name>/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 <proj_dir>
+
+# 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б.

belka_refactor_07_05_log.md

@@ -0,0 +1,208 @@
+# Шаг 4b — Что изменилось
+
+---
+
+## 1. Новая точка входа — `src/main.py` + добавлена debug-конфигурация `Main entry point` в `launch.json`
+
+**Пример запуска отладки:**
+```bash
+python src/training/train_gtauav.py --config conf/gtauav_balanced.gin
+```
+
+## 2. Тесты 
+
+Добавлены `tests\conftest.py & test_trainer.py`
+
+### Тест загрузки конфигураций в трейнер
+Файл `test_trainer.py`
+
+**Таблица конфигураций**
+
+- scr/conf/*.py конфиги
+- in/config_files общие common .gin-конфиги и пресеты под эксперименты
+---
+
+| Пресет                              | backbone  | baseline_mode | shared_encoder | init_gate |
+| ----------------------------------- | --------- | ------------- | -------------- | --------- |
+| `gtauav_balanced`                   | dinov3    | False         | True           | 0.7       |
+| `gtauav_baseline`                   | dinov3    | True          | True           | 0.7       |
+| `gtauav_balanced_asym`              | dinov3    | False         | False          | 0.7       |
+| `gtauav_baseline_asym`              | dinov3    | True          | False          | 0.7       |
+| `gtauav_text_heavy`                 | dinov3    | False         | True           | **0.3**   |
+| `gtauav_image_heavy`                | dinov3    | False         | True           | **0.9**   |
+| `gtauav_balanced_stripnet`          | stripnet  | False         | —              | 0.7       |
+| `gtauav_balanced_stripnet_unfrozen` | stripnet  | False         | —              | 0.7       |
+| `gtauav_baseline_stripnet`          | stripnet  | True          | —              | 0.7       |
+| `gtauav_baseline_stripnet_unfrozen` | stripnet  | True          | —              | 0.7       |
+| `gtauav_balanced_sofia`             | sofia_v71 | False         | —              | 0.7       |
+| `gtauav_baseline_sofia`             | sofia_v71 | True          | —              | 0.7       |
+| `gtauav_balanced_sofia_v1`          | sofia_v1  | False         | —              | 0.7       |
+| `gtauav_baseline_sofia_v1`          | sofia_v1  | True          | —              | 0.7       |
+| `preprocess`                        | —         | —             | —              | —         |
+
+
+**Пример запуска отладки:**
+```bash
+python -m pytest tests/test_trainer.py
+```
+
+**Лог:**
+
+```bash
+=========================================== test session starts ===========================================  
+platform linux -- Python 3.10.19, pytest-9.0.3, pluggy-1.6.0
+rootdir: /home/servml/Документы/code/Yaroslav/caption-test
+collected 21 items                                                                                                                                                                       
+
+tests/test_trainer.py .  ✅  _SUPPORTED_BACKBONES must be a frozenset (immutable, hashable).
+.  ✅  Exactly dinov3 and stripnet are supported in the current refactor.
+.  ✅  ModelsConfig union mirrors _SUPPORTED_BACKBONES (dinov3 | stripnet).
+.
+Supported backbones pass _validate_backbone silently.
+  ✅  dinov3
+.  ✅  stripnet
+.
+Unsupported backbones (incl. sofia) raise NotImplementedError, not ImportError.
+  ✅  sofia_v1
+.  ✅  sofia_v71
+.  ✅  mistral_42b
+.  ✅  empty
+.
+Trainer(...) instantiates from every real preset's loaded cfgs.
+  ✅  gtauav_balanced
+.  ✅  gtauav_balanced_asym
+.  ✅  gtauav_baseline
+.  ✅  gtauav_baseline_asym
+.  ✅  gtauav_image_heavy
+.  ✅  gtauav_text_heavy
+.  ✅  gtauav_balanced_stripnet
+.  ✅  gtauav_balanced_stripnet_unfrozen
+.  ✅  gtauav_baseline_stripnet
+.  ✅  gtauav_baseline_stripnet_unfrozen
+.  ✅  All runtime fields are None / 0 / [] before .train() is called.
+.  ✅  Trainer.train() takes only `self` — main.py calls trainer.train().
+
+                                                                                        [100%]
+
+===========================================  21 passed in 2.19s ===========================================
+```
+
+## 3. Декомпозиция вызовов в `trainer_new.py` 
+
+**Теперь вызовы выглядят такой цепочкой:**
+```python
+def train(self) -> None:
+        """Full pipeline: setup → build → train → evaluate → cleanup."""
+        self._validate_backbone()
+        clear_vram()
+        set_seed(self.pipeline_cfg.seed)
+        self._setup_output_dir()
+        self._setup_tracker()
+        self._build_model()
+        self._configure_gradient_checkpointing()
+        self._log_model_summary()
+        self._build_loss()
+        self._build_neg_bank()
+        self._build_data_loaders()
+        self._build_optimizer_and_scheduler()
+        self._restore_from_resume()
+        self._setup_profiler()
+
+        try:
+            self._train_loop()
+            self._final_evaluation()
+        finally:
+            self._cleanup()
+```
+
+**Лог:**
+
+### Тест вызовов до блока `try: self._train_loop()` 
+✅ Загрузка весов
+✅ Сборка пайплайна
+✅ Загрузка моделей
+✅ Загрузка данных
+
+```bash
+2026-05-07 09:40:53 caption_test.trainer INFO ⚙️  Validate backbone
+2026-05-07 09:40:53 caption_test.trainer INFO ⚙️  Setup out dir
+2026-05-07 09:40:53 caption_test.trainer INFO ⚙️ Setup tracker...
+2026-05-07 09:40:53 caption_test.trackers WARNING tensorboard not installed, skipping TB tracking
+2026-05-07 09:41:24 caption_test.trainer INFO ⚙️ Build loss...
+2026-05-07 09:41:52 caption_test.trainer INFO Building model — with text (L1/L2/L3), shared DINOv3 WEB
+2026-05-07 09:42:24 caption_test.trainer INFO ⚙️  Build DINOv3 model...
+2026-05-07 09:50:48 caption_test.model INFO 🧊 Loading DINOv3 from dinov3-vitl16-pretrain-lvd1689m.pth
+2026-05-07 09:50:48 caption_test.model INFO 🧊 Loading DINOv3 from dinov3-vitl16-pretrain-lvd1689m.pth
+2026-05-07 09:54:48 caption_test.model INFO 🧊 DINOv3 loaded: 303,129,600 params
+2026-05-07 09:54:48 caption_test.model INFO 🧊 DINOv3 loaded: 303,129,600 params
+2026-05-07 09:55:03 caption_test.adapters INFO 🔧 MONA injected: 24 adapters (blocks 12-23 of 24), 3,502,080 trainable params (bottleneck=64)
+2026-05-07 09:55:03 caption_test.adapters INFO 🔧 MONA injected: 24 adapters (blocks 12-23 of 24), 3,502,080 trainable params (bottleneck=64)
+2026-05-07 09:55:05 caption_test.model INFO Shared encoder mode: single DINOv3 WEB for drone + satellite
+2026-05-07 09:55:05 caption_test.model INFO Shared encoder mode: single DINOv3 WEB for drone + satellite
+2026-05-07 09:55:38 caption_test.dgtrs INFO 📝 Loading DGTRS-CLIP text encoder from DGTRS-CLIP-ViT-L-14.pt
+2026-05-07 09:55:38 caption_test.dgtrs INFO 📝 Loading DGTRS-CLIP text encoder from DGTRS-CLIP-ViT-L-14.pt
+2026-05-07 09:55:39 caption_test.dgtrs INFO 📝 DGTRS text encoder loaded: 123,972,096 params, context=248 tokens
+2026-05-07 09:55:39 caption_test.dgtrs INFO 📝 DGTRS text encoder loaded: 123,972,096 params, context=248 tokens
+2026-05-07 10:03:01 caption_test.adapters INFO 🔧 LoRA injected: 12 blocks, rank=4, 147,456 trainable params
+2026-05-07 10:03:01 caption_test.adapters INFO 🔧 LoRA injected: 12 blocks, rank=4, 147,456 trainable params
+2026-05-07 10:04:01 caption_test.trainer INFO embed_dim=1024
+2026-05-07 10:04:22 caption_test.trainer INFO ⚙️  Configure gradient checkpointing...
+2026-05-07 10:04:42 caption_test.trainer INFO ✅ Gradient checkpointing enabled (DINOv3 + DGTRS)
+2026-05-07 10:04:51 caption_test.trainer INFO trainable=7,059,458 (1.63%) total=434,161,154
+2026-05-07 10:04:51 caption_test.profiler INFO torchinfo not installed, using basic parameter count
+2026-05-07 10:04:51 caption_test.profiler INFO Model summary:
+Total parameters: 434,161,154
+Trainable parameters: 7,059,458 (1.63%)
+  [trainable] image_encoder.layer.12.mona_attn.gamma: [1024] (1,024)
+  ... (Вывод списка параметров)
+  ...
+2026-05-07 10:05:00 caption_test.trainer INFO ⚙️  Build loss...
+2026-05-07 10:06:06 caption_test.trainer INFO Loss: SymmetricInfoNCE  Temperature: learnable (init=0.070)  q2g=0.60 g2q=0.40
+2026-05-07 10:06:07 caption_test.trainer INFO ⚙️  Build negative bank...
+2026-05-07 10:06:13 caption_test.trainer INFO ⚙️  Build dataloaders...
+2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
+2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
+2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
+2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
+2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
+2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
+2026-05-07 10:06:19 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries                                                                                                      
+2026-05-07 10:06:19 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
+2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries                                                                                                      
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO ✅ Loaded 6252 pairs from meta/test_20.json
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO ✅ Loaded 6252 pairs from meta/test_20.json
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
+2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
+2026-05-07 10:06:24 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries                                                                                                      
+2026-05-07 10:06:24 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
+2026-05-07 10:06:25 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
+2026-05-07 10:06:25 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
+2026-05-07 10:06:25 caption_test.trainer INFO Sampler: MutuallyExclusive — no false negatives within a batch
+2026-05-07 10:06:25 caption_test.trainer INFO train=24891  test=6252  batch=8  accum=8  effective_batch=64
+2026-05-07 10:06:46 caption_test.trainer INFO ⚙️  Build optimizer & scheduler...
+2026-05-07 10:06:46 caption_test.trainer INFO Optimizer: AdamW  LR: proj=1e-04  text=1e-05  warmup=2 epochs
+2026-05-07 10:07:33 caption_test.trainer INFO ⚙️  Restore from resume...
+2026-05-07 10:07:37 caption_test.trainer INFO ⚙️  Setup profiler...
+```
+
+`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(...)` именованными аргументами
+
+

conf/balanced.gin

@@ -1,48 +0,0 @@
-# Balanced: GatedFusion with text captions enabled.
-# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs satellite
-
-import src.datasets.visloc_with_captions
-import src.losses.multi_infonce
-import src.models.dual_encoder
-import src.training.train
-
-# ---- Dual encoder ----
-DualEncoderCaptionTest.variant = "ViT-B-32"
-DualEncoderCaptionTest.pretrained_path = "checkpoints/RS5M_ViT-B-32.pt"
-DualEncoderCaptionTest.unfreeze_mode = "last_block"
-DualEncoderCaptionTest.embed_dim = 512
-DualEncoderCaptionTest.use_mlp_heads = False
-DualEncoderCaptionTest.baseline_mode = False
-DualEncoderCaptionTest.init_gate = 0.7
-DualEncoderCaptionTest.device = "cuda"
-
-# ---- Fusion ----
-GatedFusion.init_gate = 0.7
-GatedFusion.baseline_mode = False
-
-# ---- Loss ----
-InfoNCELoss.temperature_init = 0.1
-InfoNCELoss.temperature_final = 0.01
-InfoNCELoss.label_smoothing = 0.1
-InfoNCELoss.weight_q2g = 0.6
-InfoNCELoss.weight_g2q = 0.4
-
-# ---- Dataset ----
-GeoLocCaptionDataset.drop_caption_prob = 0.0
-GeoLocCaptionDataset.seed = 42
-
-# ---- Training ----
-TrainConfig.train_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/train_query.txt"
-TrainConfig.val_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/val_query.txt"
-TrainConfig.data_root = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc"
-TrainConfig.output_dir = "out/caption_test/balanced"
-TrainConfig.epochs = 10
-TrainConfig.batch_size = 128
-TrainConfig.num_workers = 4
-TrainConfig.learning_rate = 1e-4
-TrainConfig.weight_decay = 1e-4
-TrainConfig.grad_clip = 1.0
-TrainConfig.use_amp = True
-TrainConfig.eval_every = 2
-TrainConfig.seed = 42
-TrainConfig.device = "cuda"

conf/baseline_no_text.gin

@@ -1,10 +0,0 @@
-# Baseline: no text fusion (gate forced to 1.0).
-# query = drone_only -> InfoNCE vs satellite
-# Reference R@1 for delta computation.
-
-include 'balanced.gin'
-
-DualEncoderCaptionTest.baseline_mode = True
-GatedFusion.baseline_mode = True
-
-TrainConfig.output_dir = "out/caption_test/baseline_no_text"

conf/gtauav_balanced.gin

@@ -1,64 +0,0 @@
-# GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions.
-# Symmetric InfoNCE + MutuallyExclusiveSampler (no false negatives).
-# 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank.
-
-import src.losses.multi_infonce
-
-# ---- Training ----
-TrainConfigGTAUAV.epochs = 10
-TrainConfigGTAUAV.batch_size = 8
-TrainConfigGTAUAV.num_workers = 4
-TrainConfigGTAUAV.learning_rate = 1e-4
-TrainConfigGTAUAV.text_lr_factor = 0.1
-TrainConfigGTAUAV.weight_decay = 1e-4
-TrainConfigGTAUAV.grad_clip = 1.0
-TrainConfigGTAUAV.grad_accum_steps = 8
-TrainConfigGTAUAV.use_amp = True
-TrainConfigGTAUAV.eval_every = 1
-TrainConfigGTAUAV.warmup_epochs = 2
-TrainConfigGTAUAV.seed = 42
-TrainConfigGTAUAV.device = "cuda"
-
-# ---- Model ----
-TrainConfigGTAUAV.init_gate = 0.7
-TrainConfigGTAUAV.baseline_mode = False
-TrainConfigGTAUAV.shared_encoder = True  # single DINOv3 WEB for both branches
-TrainConfigGTAUAV.mona_bottleneck = 64
-TrainConfigGTAUAV.mona_last_n_blocks = 12  # inject MONA only in last 12/24 ViT blocks
-TrainConfigGTAUAV.gradient_checkpointing = True
-
-# ---- Loss ----
-TrainConfigGTAUAV.loss_type = "symmetric"
-TrainConfigGTAUAV.tau_init = 0.07
-TrainConfigGTAUAV.label_smoothing = 0.1
-TrainConfigGTAUAV.learnable_temperature = True
-TrainConfigGTAUAV.weight_q2g = 0.6
-TrainConfigGTAUAV.weight_g2q = 0.4
-TrainConfigGTAUAV.neg_bank_size = 0 # 4096
-
-# ---- Sampling ----
-TrainConfigGTAUAV.sampler_type = "mutex"  # "dss" or "mutex"
-TrainConfigGTAUAV.dss_warmup_epochs = 1  # first N epochs use mutex-only (untrained embeds not useful)
-TrainConfigGTAUAV.dss_reembed_every = 1
-TrainConfigGTAUAV.use_mutex_sampler = True  # legacy flag, kept True unless disabling both samplers
-
-# ---- Output ----
-TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"
-
-# ---- Tracking ----
-TrainConfigGTAUAV.use_wandb = False
-TrainConfigGTAUAV.use_tb = True
-TrainConfigGTAUAV.use_gradcam = False
-TrainConfigGTAUAV.gradcam_every = 5
-TrainConfigGTAUAV.use_profiler = False
-TrainConfigGTAUAV.log_grad_norms = True
-
-# ---- InfoNCELoss (gin-configurable) ----
-InfoNCELoss.temperature_init = 0.07
-InfoNCELoss.learnable_temperature = True
-InfoNCELoss.label_smoothing = 0.1
-InfoNCELoss.weight_q2g = 0.6
-InfoNCELoss.weight_g2q = 0.4
-InfoNCELoss.tau_min = 0.01
-InfoNCELoss.tau_max = 0.1
-InfoNCELoss.hard_mining_k = 0 # 512  # 0 = use whole queue (disable mining)

conf/gtauav_balanced_asym.gin

@@ -1,16 +0,0 @@
-# GTA-UAV Balanced (asymmetric, full MONA): WEB drone encoder + SAT satellite encoder.
-# MONA injected into all 24 ViT blocks of each encoder.
-# Same loss/sampling/optimizer as gtauav_balanced.gin; differs only in model arch.
-#
-# Trainable: ~17.6M (MONA 2× × 24 blocks + LoRA + TextFusionMLP + gates + tau)
-# Total params: ~733M (2× DINOv3-L + DGTRS-CLIP)
-# VRAM target (RTX 4090, 24 GB): ~16-20 GB at batch=8 with gradient checkpointing.
-
-include 'conf/gtauav_balanced.gin'
-
-# ---- Model overrides: asymmetric + full MONA ----
-TrainConfigGTAUAV.shared_encoder = False    # WEB for drone, SAT for satellite
-TrainConfigGTAUAV.mona_last_n_blocks = 24   # MONA in all 24 ViT blocks (was 12)
-
-# ---- Output ----
-TrainConfigGTAUAV.output_dir = "out/gtauav/balanced_asym"

conf/gtauav_balanced_sofia.gin

@@ -1,39 +0,0 @@
-# GTA-UAV Balanced (SOFIA-Tiny backbone): SOFIA v7.1 student trained from scratch
-# с двухуровневой text fusion:
-#   1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует feature map перед GGeM/CHP).
-#   2. Late-level: GatedFusion на дескрипторах (как в DINOv3/StripNet вариантах).
-#
-# Trainable (~5-7M):
-#   - SOFIA backbone (Tiny, ~5M, from scratch — нет pretrained)
-#   - SOFIA heads (SatHead GGeM+BN+Linear, UAVHead AltitudeFiLM+CHP+BN+Linear, +Text-FiLM)
-#   - DGTRS-CLIP LoRA (rank=4, ~147K)
-#   - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared)
-#   - Gates α_q, α_g + learnable τ
-#
-# Altitude (drone_height метры) подаётся в UAVHead.AltitudeFiLM из dataset meta CSV.
-# Для sat — altitude=None → FiLM passthrough (γ=1, β=0).
-#
-# Note: SOFIA from scratch — нужно больше эпох или warmup, чем frozen DINOv3/StripNet.
-#       Mamba-2 backend (mamba_ssm) даёт 2-8x speedup vs torch fallback.
-
-include 'conf/gtauav_balanced.gin'
-
-# ---- Backbone ----
-TrainConfigGTAUAV.backbone = "sofia"
-TrainConfigGTAUAV.sofia_preset = "Tiny"
-TrainConfigGTAUAV.sofia_d_descriptor = 1024
-TrainConfigGTAUAV.sofia_use_text_film_uav = True
-TrainConfigGTAUAV.sofia_use_text_film_sat = True
-TrainConfigGTAUAV.sofia_lora_rank = 4
-# Mamba-1 used for Tiny (Mamba-2 torch fallback has a pre-existing reshape bug
-# with channels not divisible by default headdim; switch to "mamba2" for M/L
-# presets where channels % 64 == 0 OR install mamba_ssm CUDA kernels).
-TrainConfigGTAUAV.sofia_mamba_variant = "mamba1"
-TrainConfigGTAUAV.sofia_mamba_backend = "auto"  # mamba_ssm if installed else torch fallback
-
-# ---- Training overrides ----
-TrainConfigGTAUAV.gradient_checkpointing = False  # SOFIA from-scratch — keep activations live
-TrainConfigGTAUAV.shared_encoder = True  # ignored by SOFIA but kept for logging compat
-
-# ---- Output ----
-TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia"

conf/gtauav_balanced_sofia_v1.gin

@@ -1,42 +0,0 @@
-# GTA-UAV Balanced (SOFIA v1 backbone): StripNet+DCNv4 hierarchical CNN
-# (~3-30M params depending on variant) trained from scratch с двухуровневой
-# text fusion:
-#   1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует [B,C,8,8] перед GGeM).
-#   2. Late-level: GatedFusion на дескрипторах.
-#
-# UAV head: AltitudeFiLM(drone_height) + [TextFiLM] + GGeM + BN + Linear.
-# SAT head: [TextFiLM] + GGeM + BN + Linear.
-# Один backbone shared между sat/uav.
-#
-# Variant -> размер модели:
-#   tiny_tiny: dims [16, 32, 80, 128]   (~0.4M)
-#   tiny     : dims [32, 64, 128, 256]  (~1M)
-#   small    : dims [64, 128, 320, 512] (~3M, default)
-#   small_v2 : dims [64, 128, 256, 384] (~2M)
-#
-# Trainable (с small variant + text fusion):
-#   - SOFIA v1 backbone (~3M) + heads (~0.6M)
-#   - DGTRS-CLIP LoRA (rank=4, ~147K)
-#   - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared)
-#   - Gates α_q, α_g + learnable τ
-# Total trainable ~7M.
-#
-# Note: DCNv4 требует CUDA — обучение только на GPU. Не работает на CPU.
-
-include 'conf/gtauav_balanced.gin'
-
-# ---- Backbone ----
-TrainConfigGTAUAV.backbone = "sofia_v1"
-TrainConfigGTAUAV.sofia_v1_variant = "tiny"
-TrainConfigGTAUAV.sofia_v1_d_descriptor = 1024
-TrainConfigGTAUAV.sofia_v1_use_text_film_uav = True
-TrainConfigGTAUAV.sofia_v1_use_text_film_sat = True
-TrainConfigGTAUAV.sofia_v1_use_film_altitude = True
-TrainConfigGTAUAV.sofia_v1_lora_rank = 4
-
-# ---- Training overrides ----
-TrainConfigGTAUAV.gradient_checkpointing = False  # SOFIA v1 from-scratch
-TrainConfigGTAUAV.shared_encoder = True
-
-# ---- Output ----
-TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia_v1"

conf/gtauav_balanced_stripnet.gin

@@ -1,27 +0,0 @@
-# GTA-UAV Balanced (StripNet backbone): StripNet-small + Conv-MONA in last 2 stages.
-# Replaces DINOv3 ViT-L/16 with strip-shaped DWConv hierarchical CNN (~28M params,
-# 10× smaller than DINOv3). Output 512-dim → projected to 1024 to match retrieval space.
-#
-# Trainable:
-#   - Projection (Linear 512→1024): ~525K
-#   - Conv-MONA in stages 3 & 4 (2 adapters per Block × 6 blocks total): ~2-3M
-#   - LoRA on DGTRS-CLIP: 147K
-#   - TextFusionMLP (shared): ~3.4M
-#   - GatedFusion gates + tau: 3 scalars
-#
-# StripNet pretrained on ImageNet-1K (head dropped); state-dict naming follows
-# upstream Strip-R-CNN repo (`conv_spatial1/2`).
-
-include 'conf/gtauav_balanced.gin'
-
-# ---- Backbone ----
-TrainConfigGTAUAV.backbone = "stripnet"
-TrainConfigGTAUAV.stripnet_path = "nn_models/STRIPNET/stripnet_s.pth"
-TrainConfigGTAUAV.stripnet_mona_last_n_stages = 2  # Conv-MONA in stages 3 & 4 (deepest)
-
-# ---- Model overrides ----
-TrainConfigGTAUAV.shared_encoder = True   # StripNet always shared (one CNN for both branches)
-TrainConfigGTAUAV.mona_bottleneck = 64    # Conv-MONA bottleneck channels
-
-# ---- Output ----
-TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_stripnet"

conf/gtauav_balanced_stripnet_unfrozen.gin

@@ -1,17 +0,0 @@
-# GTA-UAV Balanced (StripNet, fully unfrozen): all StripNet layers trainable.
-# Backbone trains with reduced LR (lr * stripnet_backbone_lr_factor).
-# Conv-MONA disabled by default — full fine-tune supplies enough capacity.
-# Set stripnet_mona_last_n_stages > 0 if you want MONA + fine-tune hybrid.
-#
-# Note: StripNet uses BatchNorm. With small batch (8) and gradient accumulation,
-# BN running stats may drift. Watch validation loss for instability.
-
-include 'conf/gtauav_balanced_stripnet.gin'
-
-# ---- Unfreeze backbone ----
-TrainConfigGTAUAV.stripnet_freeze = False
-TrainConfigGTAUAV.stripnet_mona_last_n_stages = 0  # disable Conv-MONA (full fine-tune handles adaptation)
-TrainConfigGTAUAV.stripnet_backbone_lr_factor = 0.1  # backbone lr = 1e-4 * 0.1 = 1e-5
-
-# ---- Output ----
-TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_stripnet_unfrozen"

conf/gtauav_baseline.gin

@@ -1,28 +0,0 @@
-# GTA-UAV Baseline: no text fusion (gate forced to 1.0).
-# query = drone_only -> InfoNCE vs satellite
-# Reference R@1 for delta computation.
-#
-# Diagnostic mode (2026-04-24): DSS and hard-negative mining disabled after
-# the previous run collapsed (R@1 = 0.6% at epoch 8, train loss growing).
-# Hypothesis: DSS packs visually-identical drones at bs=8 and the hard-mining
-# queue amplifies that hardness — together they prevent convergence from a
-# nearly-random start. Run with mutex-only sampling and the full queue as
-# uniform negatives first, restore the extras incrementally once baseline
-# converges.
-
-include 'conf/gtauav_balanced.gin'
-
-TrainConfigGTAUAV.baseline_mode = True
-TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_inbatch"
-TrainConfigGTAUAV.use_gradcam = False
-
-# ---- Diagnostic overrides ----
-# Previous mutex-only run still collapsed at epoch 1 (val loss locked at log(8)).
-# Hypothesis refined: the MoCo-style queue stays stale because we have no
-# momentum encoder, and with reduced trainable surface (MONA-12) the model
-# can't reconcile fresh representations against 4096 stale negatives —
-# mode collapse. Disable the queue entirely so InfoNCE sees only the 8
-# fresh in-batch negatives, matching the OLD run's effective setup.
-TrainConfigGTAUAV.sampler_type = "mutex"  # was "dss"
-TrainConfigGTAUAV.neg_bank_size = 0  # was 4096 — disable MoCo queue (no momentum encoder)
-InfoNCELoss.hard_mining_k = 0  # was 512 — irrelevant when queue is empty

conf/gtauav_baseline_asym.gin

@@ -1,9 +0,0 @@
-# GTA-UAV Baseline (asymmetric, full MONA): no text fusion (gate forced to 1.0).
-# WEB drone encoder + SAT satellite encoder, MONA in all 24 ViT blocks.
-# Reference R@1 for delta computation against gtauav_balanced_asym.gin.
-
-include 'conf/gtauav_balanced_asym.gin'
-
-TrainConfigGTAUAV.baseline_mode = True
-TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_asym"
-TrainConfigGTAUAV.use_gradcam = False

conf/gtauav_baseline_sofia.gin

@@ -1,13 +0,0 @@
-# GTA-UAV Baseline (SOFIA-Tiny backbone): no text fusion. Reference R@1 для
-# computing Δ R@1 vs gtauav_balanced_sofia.gin.
-#
-# В baseline_mode=True:
-#   - Text-FiLM отключается (SOFIA heads работают только с altitude).
-#   - DGTRS-CLIP не загружается, TextFusionMLP не строится.
-#   - GatedFusion gates = 1.0 (text игнорируется).
-
-include 'conf/gtauav_balanced_sofia.gin'
-
-TrainConfigGTAUAV.baseline_mode = True
-TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia"
-TrainConfigGTAUAV.use_gradcam = False

conf/gtauav_baseline_sofia_v1.gin

@@ -1,8 +0,0 @@
-# GTA-UAV Baseline (SOFIA v1 backbone): no text fusion. Reference R@1 для
-# computing Δ R@1 vs gtauav_balanced_sofia_v1.gin.
-
-include 'conf/gtauav_balanced_sofia_v1.gin'
-
-TrainConfigGTAUAV.baseline_mode = True
-TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia_v1"
-TrainConfigGTAUAV.use_gradcam = False

conf/gtauav_baseline_stripnet.gin

@@ -1,8 +0,0 @@
-# GTA-UAV Baseline (StripNet backbone): no text fusion. Reference R@1 for
-# computing Δ R@1 against gtauav_balanced_stripnet.gin.
-
-include 'conf/gtauav_balanced_stripnet.gin'
-
-TrainConfigGTAUAV.baseline_mode = True
-TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_stripnet"
-TrainConfigGTAUAV.use_gradcam = False

conf/gtauav_baseline_stripnet_unfrozen.gin

@@ -1,8 +0,0 @@
-# GTA-UAV Baseline (StripNet, fully unfrozen): no text fusion. For Δ R@1
-# against gtauav_balanced_stripnet_unfrozen.gin.
-
-include 'conf/gtauav_balanced_stripnet_unfrozen.gin'
-
-TrainConfigGTAUAV.baseline_mode = True
-TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_stripnet_unfrozen"
-TrainConfigGTAUAV.use_gradcam = False

conf/gtauav_image_heavy.gin

@@ -1,8 +0,0 @@
-# GTA-UAV Image-heavy: gate initialized high (more image weight).
-# query = sigma(0.9) * drone + 0.1 * text
-# Minimal text contribution test.
-
-include 'conf/gtauav_balanced.gin'
-
-TrainConfigGTAUAV.init_gate = 0.9
-TrainConfigGTAUAV.output_dir = "out/gtauav/image_heavy"

conf/gtauav_text_heavy.gin

@@ -1,8 +0,0 @@
-# GTA-UAV Text-heavy: gate initialized low (more text weight).
-# query = sigma(0.3) * drone + 0.7 * text
-# Stress test for text contribution.
-
-include 'conf/gtauav_balanced.gin'
-
-TrainConfigGTAUAV.init_gate = 0.3
-TrainConfigGTAUAV.output_dir = "out/gtauav/text_heavy"

conf/text_heavy.gin

@@ -1,9 +0,0 @@
-# Text-heavy: gate initialized low (more text weight).
-# query = sigma(0.3) * drone + 0.7 * text
-
-include 'balanced.gin'
-
-DualEncoderCaptionTest.init_gate = 0.3
-GatedFusion.init_gate = 0.3
-
-TrainConfig.output_dir = "out/caption_test/text_heavy"

in/config_files/gtauav_balanced/models.gin

@@ -0,0 +1,12 @@
+# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
+ModelsCommonConfig.backbone = 'dinov3'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
+DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
+DINOv3ModelsConfig.shared_encoder = True
+DINOv3ModelsConfig.mona_bottleneck = 64
+DINOv3ModelsConfig.mona_last_n_blocks = 12
+DINOv3ModelsConfig.lora_rank = 4

in/config_files/gtauav_balanced/pipeline.gin

@@ -0,0 +1,14 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/with_text'
+PipelineConfig.resume_from = None

in/config_files/gtauav_balanced_asym/models.gin

@@ -0,0 +1,11 @@
+ModelsCommonConfig.backbone = 'dinov3'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
+DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
+DINOv3ModelsConfig.shared_encoder = False            
+DINOv3ModelsConfig.mona_bottleneck = 64
+DINOv3ModelsConfig.mona_last_n_blocks = 24           
+DINOv3ModelsConfig.lora_rank = 4

in/config_files/gtauav_balanced_asym/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/balanced_asym'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_balanced_sofia/models.gin

@@ -0,0 +1,91 @@
+# SOFIA v7.1 Tiny preset (~5M params) with text fusion (Text-FiLM mid-level
+# in SAT and UAV heads + GatedFusion late-level on descriptors).
+#
+# Tiny-specific notes:
+#   - num_heads_s3/s4 = 4 (channels 176/224 not divisible by 8)
+#   - mamba_headdim = 16 (channels not divisible by default 64)
+#   - mamba_variant = 'mamba1' (Mamba-2 torch fallback bug for these dims)
+#   - d_descriptor = 1024 (override from preset M default 512)
+#   - text fusion enabled (override from preset M default disabled)
+
+ModelsCommonConfig.backbone = 'sofia_v71'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+# Variant label (informational).
+SOFIAv71ModelsConfig.variant_label = 'Tiny'
+
+# Input.
+SOFIAv71ModelsConfig.input_size = 256
+SOFIAv71ModelsConfig.in_channels = 3
+
+# Stem (Tiny dims).
+SOFIAv71ModelsConfig.stem_mid = 16
+SOFIAv71ModelsConfig.stem_out = 32
+
+# Backbone dimensions (Tiny).
+SOFIAv71ModelsConfig.embed_dims = [48, 96, 176, 224]
+SOFIAv71ModelsConfig.depths = [2, 3, 4, 2]
+
+# Stage 1-2 block params (default).
+SOFIAv71ModelsConfig.mbconv_expand = 4
+SOFIAv71ModelsConfig.se_ratio = 16
+SOFIAv71ModelsConfig.strip_kernel_s1 = 7
+SOFIAv71ModelsConfig.strip_kernel_s2 = 5
+SOFIAv71ModelsConfig.mix_kernels = [3, 5, 7]
+SOFIAv71ModelsConfig.use_dcn_strip = True
+
+# Stage 3-4 (MambaVision). Tiny: mamba1 to bypass torch fallback bug.
+SOFIAv71ModelsConfig.mamba_d_state = 16
+SOFIAv71ModelsConfig.mamba_dt_rank = None
+SOFIAv71ModelsConfig.mamba_backend = 'auto'
+SOFIAv71ModelsConfig.mamba_variant = 'mamba1'
+
+# Mamba-2 tunables (used when mamba_variant='mamba2'; Tiny would need
+# headdim=16 because 176 % 64 != 0 and 224 % 64 != 0).
+SOFIAv71ModelsConfig.mamba_d_state_mamba2 = 64
+SOFIAv71ModelsConfig.mamba_headdim = 16
+SOFIAv71ModelsConfig.mamba_expand = 2
+SOFIAv71ModelsConfig.mamba_d_conv = 4
+SOFIAv71ModelsConfig.mamba_n_directions = 2
+
+# Heads / attention (Tiny: heads=4).
+SOFIAv71ModelsConfig.num_heads_s3 = 4
+SOFIAv71ModelsConfig.num_heads_s4 = 4
+SOFIAv71ModelsConfig.use_strip_branch_s3 = True
+SOFIAv71ModelsConfig.use_strip_branch_s4 = False
+SOFIAv71ModelsConfig.ffn_expand = 4
+
+# EVSS bridge (off by default).
+SOFIAv71ModelsConfig.use_evss_bridge = False
+SOFIAv71ModelsConfig.evss_bridge_locations = ['pre_stage3']
+
+# Neck (Tiny).
+SOFIAv71ModelsConfig.neck_channels = 128
+
+# CVGL Head.
+SOFIAv71ModelsConfig.d_descriptor = 1024
+SOFIAv71ModelsConfig.use_asymmetric_heads = True
+SOFIAv71ModelsConfig.chp_rings = 8
+SOFIAv71ModelsConfig.chp_angles = 16
+SOFIAv71ModelsConfig.chp_harmonics = 4
+SOFIAv71ModelsConfig.use_film_altitude = True
+SOFIAv71ModelsConfig.altitude_norm = 500.0
+SOFIAv71ModelsConfig.ring_count = 4
+SOFIAv71ModelsConfig.use_ring_aux = True
+
+# Text fusion enabled.
+SOFIAv71ModelsConfig.return_normalized = True
+SOFIAv71ModelsConfig.use_text_film_sat = True
+SOFIAv71ModelsConfig.use_text_film_uav = True
+SOFIAv71ModelsConfig.text_film_dim = 1024
+SOFIAv71ModelsConfig.text_film_hidden = 256
+
+# Sharing / KD / deploy.
+SOFIAv71ModelsConfig.share_stages_1_2 = True
+SOFIAv71ModelsConfig.enable_kd_taps = True
+SOFIAv71ModelsConfig.precision = 'fp16'
+
+# LoRA.
+SOFIAv71ModelsConfig.lora_rank = 4

in/config_files/gtauav_balanced_sofia/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/with_text_sofia'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_balanced_sofia_v1/models.gin

@@ -0,0 +1,30 @@
+# SOFIA v1 'tiny' variant (~1M params) with text fusion (Text-FiLM mid-level
+# in SAT and UAV heads + AltitudeFiLM in UAV head + GatedFusion late-level).
+
+ModelsCommonConfig.backbone = 'sofia_v1'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+# Backbone.
+SOFIAv1ModelsConfig.variant_label = 'tiny'
+SOFIAv1ModelsConfig.in_channels = 3
+SOFIAv1ModelsConfig.input_size = 256
+SOFIAv1ModelsConfig.dcn_variant = 'v2'
+
+# Heads.
+SOFIAv1ModelsConfig.d_descriptor = 1024
+SOFIAv1ModelsConfig.return_normalized = False
+
+# Altitude-FiLM.
+SOFIAv1ModelsConfig.use_film_altitude = True
+SOFIAv1ModelsConfig.altitude_norm = 500.0
+
+# Text-FiLM.
+SOFIAv1ModelsConfig.use_text_film_uav = True
+SOFIAv1ModelsConfig.use_text_film_sat = True
+SOFIAv1ModelsConfig.text_film_dim = 1024
+SOFIAv1ModelsConfig.text_film_hidden = 256
+
+# LoRA on DGTRS-CLIP.
+SOFIAv1ModelsConfig.lora_rank = 4

in/config_files/gtauav_balanced_sofia_v1/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/with_text_sofia_v1'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_balanced_stripnet/models.gin

@@ -0,0 +1,11 @@
+# StripNet small backbone (frozen) + Conv-MONA on last 2 stages.
+ModelsCommonConfig.backbone = 'stripnet'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
+StripNetModelsConfig.stripnet_freeze = True
+StripNetModelsConfig.stripnet_mona_last_n_stages = 2
+StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
+StripNetModelsConfig.lora_rank = 4

in/config_files/gtauav_balanced_stripnet/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/balanced_stripnet'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_balanced_stripnet_unfrozen/models.gin

@@ -0,0 +1,11 @@
+# StripNet small backbone (frozen) + Conv-MONA on last 2 stages.
+ModelsCommonConfig.backbone = 'stripnet'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
+StripNetModelsConfig.stripnet_freeze = False
+StripNetModelsConfig.stripnet_mona_last_n_stages = 2
+StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
+StripNetModelsConfig.lora_rank = 4

in/config_files/gtauav_balanced_stripnet_unfrozen/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/balanced_stripnet_unfrozen'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_baseline/models.gin

@@ -0,0 +1,12 @@
+# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
+ModelsCommonConfig.backbone = 'dinov3'
+ModelsCommonConfig.baseline_mode = True
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
+DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
+DINOv3ModelsConfig.shared_encoder = True
+DINOv3ModelsConfig.mona_bottleneck = 64
+DINOv3ModelsConfig.mona_last_n_blocks = 12
+DINOv3ModelsConfig.lora_rank = 4

in/config_files/gtauav_baseline/pipeline.gin

@@ -0,0 +1,14 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/baseline_inbatch'
+PipelineConfig.resume_from = None

in/config_files/gtauav_baseline_asym/models.gin

@@ -0,0 +1,12 @@
+# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
+ModelsCommonConfig.backbone = 'dinov3'
+ModelsCommonConfig.baseline_mode = True
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
+DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
+DINOv3ModelsConfig.shared_encoder = False
+DINOv3ModelsConfig.mona_bottleneck = 64
+DINOv3ModelsConfig.mona_last_n_blocks = 24
+DINOv3ModelsConfig.lora_rank = 4

in/config_files/gtauav_baseline_asym/pipeline.gin

@@ -0,0 +1,14 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/baseline_asym'
+PipelineConfig.resume_from = None

in/config_files/gtauav_baseline_sofia/models.gin

@@ -0,0 +1,92 @@
+# SOFIA v7.1 Tiny preset (~5M params) with text fusion (Text-FiLM mid-level
+# in SAT and UAV heads + GatedFusion late-level on descriptors).
+#
+# Tiny-specific notes:
+#   - num_heads_s3/s4 = 4 (channels 176/224 not divisible by 8)
+#   - mamba_headdim = 16 (channels not divisible by default 64)
+#   - mamba_variant = 'mamba1' (Mamba-2 torch fallback bug for these dims)
+#   - d_descriptor = 1024 (override from preset M default 512)
+#   - text fusion enabled (override from preset M default disabled)
+
+ModelsCommonConfig.backbone = 'sofia_v71'
+ModelsCommonConfig.baseline_mode = True
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+# Variant label (informational).
+SOFIAv71ModelsConfig.variant_label = 'Tiny'
+
+# Input.
+SOFIAv71ModelsConfig.input_size = 256
+SOFIAv71ModelsConfig.in_channels = 3
+
+# Stem (Tiny dims).
+SOFIAv71ModelsConfig.stem_mid = 16
+SOFIAv71ModelsConfig.stem_out = 32
+
+# Backbone dimensions (Tiny).
+SOFIAv71ModelsConfig.embed_dims = [48, 96, 176, 224]
+SOFIAv71ModelsConfig.depths = [2, 3, 4, 2]
+
+# Stage 1-2 block params (default).
+SOFIAv71ModelsConfig.mbconv_expand = 4
+SOFIAv71ModelsConfig.se_ratio = 16
+SOFIAv71ModelsConfig.strip_kernel_s1 = 7
+SOFIAv71ModelsConfig.strip_kernel_s2 = 5
+SOFIAv71ModelsConfig.mix_kernels = [3, 5, 7]
+SOFIAv71ModelsConfig.use_dcn_strip = True
+
+# Stage 3-4 (MambaVision). Tiny: mamba1 to bypass torch fallback bug.
+SOFIAv71ModelsConfig.mamba_d_state = 16
+SOFIAv71ModelsConfig.mamba_dt_rank = None
+SOFIAv71ModelsConfig.mamba_backend = 'auto'
+SOFIAv71ModelsConfig.mamba_variant = 'mamba1'
+
+# Mamba-2 tunables (used when mamba_variant='mamba2'; Tiny would need
+# headdim=16 because 176 % 64 != 0 and 224 % 64 != 0).
+SOFIAv71ModelsConfig.mamba_d_state_mamba2 = 64
+SOFIAv71ModelsConfig.mamba_headdim = 16
+SOFIAv71ModelsConfig.mamba_expand = 2
+SOFIAv71ModelsConfig.mamba_d_conv = 4
+SOFIAv71ModelsConfig.mamba_n_directions = 2
+
+# Heads / attention (Tiny: heads=4).
+SOFIAv71ModelsConfig.num_heads_s3 = 4
+SOFIAv71ModelsConfig.num_heads_s4 = 4
+SOFIAv71ModelsConfig.use_strip_branch_s3 = True
+SOFIAv71ModelsConfig.use_strip_branch_s4 = False
+SOFIAv71ModelsConfig.ffn_expand = 4
+
+# EVSS bridge (off by default).
+SOFIAv71ModelsConfig.use_evss_bridge = False
+SOFIAv71ModelsConfig.evss_bridge_locations = ['pre_stage3']
+
+# Neck (Tiny).
+SOFIAv71ModelsConfig.neck_channels = 128
+
+# CVGL Head.
+SOFIAv71ModelsConfig.d_descriptor = 1024
+SOFIAv71ModelsConfig.use_asymmetric_heads = True
+SOFIAv71ModelsConfig.chp_rings = 8
+SOFIAv71ModelsConfig.chp_angles = 16
+SOFIAv71ModelsConfig.chp_harmonics = 4
+SOFIAv71ModelsConfig.use_film_altitude = True
+SOFIAv71ModelsConfig.altitude_norm = 500.0
+SOFIAv71ModelsConfig.ring_count = 4
+SOFIAv71ModelsConfig.use_ring_aux = True
+
+# Text fusion enabled.
+SOFIAv71ModelsConfig.return_normalized = True
+SOFIAv71ModelsConfig.use_text_film_sat = True
+SOFIAv71ModelsConfig.use_text_film_uav = True
+SOFIAv71ModelsConfig.text_film_dim = 1024
+SOFIAv71ModelsConfig.text_film_hidden = 256
+
+# Sharing / KD / deploy.
+SOFIAv71ModelsConfig.share_stages_1_2 = True
+SOFIAv71ModelsConfig.enable_kd_taps = True
+SOFIAv71ModelsConfig.precision = 'fp16'
+
+# LoRA.
+SOFIAv71ModelsConfig.lora_rank = 4
+

in/config_files/gtauav_baseline_sofia/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/baseline_sofia'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_baseline_sofia_v1/models.gin

@@ -0,0 +1,31 @@
+# SOFIA v1 'tiny' variant (~1M params) with text fusion (Text-FiLM mid-level
+# in SAT and UAV heads + AltitudeFiLM in UAV head + GatedFusion late-level).
+
+ModelsCommonConfig.backbone = 'sofia_v1'
+ModelsCommonConfig.baseline_mode = True
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+# Backbone.
+SOFIAv1ModelsConfig.variant_label = 'tiny'
+SOFIAv1ModelsConfig.in_channels = 3
+SOFIAv1ModelsConfig.input_size = 256
+SOFIAv1ModelsConfig.dcn_variant = 'v2'
+
+# Heads.
+SOFIAv1ModelsConfig.d_descriptor = 1024
+SOFIAv1ModelsConfig.return_normalized = False
+
+# Altitude-FiLM.
+SOFIAv1ModelsConfig.use_film_altitude = True
+SOFIAv1ModelsConfig.altitude_norm = 500.0
+
+# Text-FiLM.
+SOFIAv1ModelsConfig.use_text_film_uav = True
+SOFIAv1ModelsConfig.use_text_film_sat = True
+SOFIAv1ModelsConfig.text_film_dim = 1024
+SOFIAv1ModelsConfig.text_film_hidden = 256
+
+# LoRA on DGTRS-CLIP.
+SOFIAv1ModelsConfig.lora_rank = 4
+

in/config_files/gtauav_baseline_sofia_v1/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/baseline_sofia_v1'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_baseline_stripnet/models.gin

@@ -0,0 +1,12 @@
+# StripNet small backbone (frozen) + Conv-MONA on last 2 stages.
+ModelsCommonConfig.backbone = 'stripnet'
+ModelsCommonConfig.baseline_mode = True
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
+StripNetModelsConfig.stripnet_freeze = True
+StripNetModelsConfig.stripnet_mona_last_n_stages = 2
+StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
+StripNetModelsConfig.lora_rank = 4
+

in/config_files/gtauav_baseline_stripnet/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/baseline_stripnet'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_baseline_stripnet_unfrozen/models.gin

@@ -0,0 +1,11 @@
+# StripNet small backbone (unfrozen) + Conv-MONA on last 2 stages.
+ModelsCommonConfig.backbone = 'stripnet'
+ModelsCommonConfig.baseline_mode = True
+ModelsCommonConfig.init_gate = 0.7
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
+StripNetModelsConfig.stripnet_freeze = False
+StripNetModelsConfig.stripnet_mona_last_n_stages = 2
+StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
+StripNetModelsConfig.lora_rank = 4

in/config_files/gtauav_baseline_stripnet_unfrozen/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/baseline_stripnet_unfrozen'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_image_heavy/models.gin

@@ -0,0 +1,13 @@
+# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
+ModelsCommonConfig.backbone = 'dinov3'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.9
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
+DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
+DINOv3ModelsConfig.shared_encoder = True
+DINOv3ModelsConfig.mona_bottleneck = 64
+DINOv3ModelsConfig.mona_last_n_blocks = 12
+DINOv3ModelsConfig.lora_rank = 4
+

in/config_files/gtauav_image_heavy/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/image_heavy'
+PipelineConfig.resume_from = None
+

in/config_files/gtauav_text_heavy/models.gin

@@ -0,0 +1,13 @@
+# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
+ModelsCommonConfig.backbone = 'dinov3'
+ModelsCommonConfig.baseline_mode = False
+ModelsCommonConfig.init_gate = 0.3
+ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
+
+DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
+DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
+DINOv3ModelsConfig.shared_encoder = True
+DINOv3ModelsConfig.mona_bottleneck = 64
+DINOv3ModelsConfig.mona_last_n_blocks = 12
+DINOv3ModelsConfig.lora_rank = 4
+

in/config_files/gtauav_text_heavy/pipeline.gin

@@ -0,0 +1,15 @@
+# Pipeline: GTA-UAV-LR with text captions, servml workstation paths.
+PipelineConfig.train_json = 'meta/train_80.json'
+PipelineConfig.test_json = 'meta/test_20.json'
+PipelineConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions'
+PipelineConfig.filter_meta = 'meta/seg_filter.json'
+
+PipelineConfig.epochs = 10
+PipelineConfig.warmup_epochs = 2
+PipelineConfig.eval_every = 1
+PipelineConfig.seed = 42
+
+PipelineConfig.output_dir = 'out/gtauav/text_heavy'
+PipelineConfig.resume_from = None
+

in/config_files/hardware_default.gin

@@ -0,0 +1,9 @@
+# RTX 4090 profile, shared encoder (DINOv3).
+HardwareConfig.device = 'cuda'
+HardwareConfig.batch_size = 8
+HardwareConfig.grad_accum_steps = 8
+HardwareConfig.num_workers = 4
+HardwareConfig.use_amp = True
+HardwareConfig.gradient_checkpointing = True
+HardwareConfig.reserve_gb = 2.0
+HardwareConfig.total_vram_gb = 24.0

in/config_files/hardware_no_gradckpt.gin

@@ -0,0 +1,10 @@
+# SOFIA v7.1 from-scratch — keep activations live (no gradient checkpointing).
+HardwareConfig.device = 'cuda'
+HardwareConfig.batch_size = 8
+HardwareConfig.grad_accum_steps = 8
+HardwareConfig.num_workers = 4
+HardwareConfig.use_amp = True
+HardwareConfig.gradient_checkpointing = False
+HardwareConfig.reserve_gb = 2.0
+HardwareConfig.total_vram_gb = 24.0
+

in/config_files/preprocess/preprocess.gin

@@ -0,0 +1,20 @@
+# Preprocessing config used by scripts/make_split.py and
+# scripts/filter_segmentation.py. Independent from training pipeline.
+
+# Inputs.
+PreprocessConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
+PreprocessConfig.segm_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm'
+
+# make_split.py — 80/20 split with seed=42.
+PreprocessConfig.split_ratio = 0.8
+PreprocessConfig.split_seed = 42
+PreprocessConfig.split_input_train = 'cross-area-drone2sate-train.json'
+PreprocessConfig.split_input_test = 'cross-area-drone2sate-test.json'
+PreprocessConfig.split_output_dir = 'meta'
+PreprocessConfig.split_output_train = 'train_80.json'
+PreprocessConfig.split_output_test = 'test_20.json'
+
+# filter_segmentation.py — exclude images with >=90% background+water.
+PreprocessConfig.seg_threshold = 0.90
+PreprocessConfig.seg_exclude_classes = [0, 4]
+PreprocessConfig.seg_filter_output = 'meta/seg_filter.json'

in/config_files/tracking.gin

@@ -0,0 +1,16 @@
+# Default tracking: TensorBoard on, W&B off, no Grad-CAM, no profiler.
+TrackingConfig.use_wandb = False
+TrackingConfig.use_tb = True
+TrackingConfig.wandb_project = 'caption-test-gtauav'
+TrackingConfig.wandb_run_name = None
+TrackingConfig.wandb_entity = None
+
+TrackingConfig.log_grad_norms = True
+TrackingConfig.use_gradcam = False
+TrackingConfig.gradcam_every = 5
+TrackingConfig.gradcam_samples = 8
+
+TrackingConfig.use_profiler = False
+TrackingConfig.profiler_warmup = 3
+TrackingConfig.profiler_active = 5
+

in/config_files/training.gin

@@ -0,0 +1,31 @@
+# Loss + optimizer + sampler — symmetric InfoNCE, AdamW, mutex sampler.
+TrainingConfig.loss_type = 'symmetric'
+TrainingConfig.tau_init = 0.07
+TrainingConfig.tau_min = 0.01
+TrainingConfig.tau_max = 0.1
+TrainingConfig.learnable_temperature = True
+TrainingConfig.label_smoothing = 0.1
+
+TrainingConfig.tau_final = 0.01
+TrainingConfig.weight_q2g = 0.6
+TrainingConfig.weight_g2q = 0.4
+TrainingConfig.hard_mining_k = 0
+TrainingConfig.neg_bank_size = 0
+
+# WeightedInfoNCE-only (unused when loss_type='symmetric').
+TrainingConfig.weighted_loss_k = 5.0
+
+TrainingConfig.learning_rate = 1e-4
+TrainingConfig.text_lr_factor = 0.1
+TrainingConfig.weight_decay = 1e-4
+TrainingConfig.grad_clip = 1.0
+
+TrainingConfig.sampler_type = 'mutex'
+TrainingConfig.dss_warmup_epochs = 1
+TrainingConfig.dss_reembed_every = 1
+TrainingConfig.dss_knn_device = 'cuda'
+TrainingConfig.dss_use_lsh = False
+TrainingConfig.dss_lsh_num_tables = 8
+TrainingConfig.dss_lsh_num_bits = 14
+TrainingConfig.dss_cache_dir = None
+TrainingConfig.use_mutex_sampler = True

pre_refactor_analysis.md

@@ -0,0 +1,297 @@
+# Шаг 1 — Что не соответствует стандарту в текущем состоянии
+
+> **Только диагностика.** Никаких решений по рефакторингу — это пойдёт на следующих шагах. Здесь только: что есть сейчас, какое правило стандарта нарушено, и каков масштаб проблемы.
+
+---
+
+## Текущее состояние
+
+| Что | Где | Сколько |
+|---|---|---|
+| Гин-конфиги (взаимовключаемые через `include`) | `conf/` | 12 файлов: `balanced.gin`, `baseline_no_text.gin`, `text_heavy.gin` (v2 UAV-GeoLoc), `gtauav_balanced.gin`, `gtauav_baseline.gin`, `gtauav_balanced_asym.gin`, `gtauav_baseline_asym.gin`, `gtauav_balanced_stripnet.gin`, `gtauav_balanced_stripnet_unfrozen.gin`, `gtauav_baseline_stripnet.gin`, `gtauav_baseline_stripnet_unfrozen.gin`, `gtauav_text_heavy.gin`, `gtauav_image_heavy.gin` (по скриншоту в локальной копии есть ещё 4 sofia-варианта) |
+| Главный `@gin.configurable + @dataclass` | `src/training/train_gtauav.py::TrainConfigGTAUAV` | 1 класс, 50+ полей |
+| Legacy `@gin.configurable` на классе тренировки | `src/training/train.py::TrainConfig` (v2) | 1 класс, 14 полей |
+| `@gin.configurable` на не-конфиг классах | `src/losses/multi_infonce.py::InfoNCELoss`, `src/losses/weighted_infonce.py::WeightedInfoNCELoss`, `src/datasets/visloc_with_captions.py::GeoLocCaptionDataset` | 3 класса |
+| Локальный `@dataclass` для модели (без gin) | `src/models/sofia_v1/config.py::SOFIAv1Config` | 1 класс |
+| Module-level пути | `src/training/train_gtauav.py`, `src/datasets/gtauav_dataset.py`, `scripts/make_split.py`, `scripts/filter_segmentation.py` | 4 файла |
+| `argparse` | `src/training/train_gtauav.py`, `src/training/train.py`, `scripts/make_split.py`, `scripts/filter_segmentation.py` | 4 файла |
+
+---
+
+## Несоответствия стандарту
+
+### 1. `@gin.configurable + @dataclass` — критическое нарушение
+
+**Правило** (`Стандарт_написания_кода_для_DL_CV.md` §3.1, `Reference Examples → Anti-patterns`):
+
+> «**Запрещено** использовать `dataclass` совместно с gin»
+
+**Что нарушает:**
+
+`src/training/train_gtauav.py`:
+```python
+@gin.configurable(module="src.training.train_gtauav")
+@dataclass                                              # ← FORBIDDEN
+class TrainConfigGTAUAV:
+    train_json: str = _TRAIN_JSON
+    test_json: str = _TEST_JSON
+    rgb_root: str = _RGB_ROOT
+    # ... 50+ полей всего:
+    # пути к данным, к моделям, training schedule, model arch (mona, gates,
+    # baseline, stripnet, asymmetric), loss params, sampler params,
+    # tracking flags (wandb/tb/gradcam/profiler)
+```
+
+**Почему это нарушение:** `@dataclass` авто-генерирует `__init__`, в который gin **тоже** инъектирует параметры. Получается двойная магия: dataclass читает type hints для генерации сигнатуры, gin читает гин-биндинги для подмены значений. На практике работает, но это именно та комбинация, которую стандарт запрещает.
+
+**Масштаб:** **1 класс**, но он держит **всё** — путь, hardware, model, loss, sampler, tracking. Это блокирует разделение конфига на оси.
+
+---
+
+### 2. `@gin.configurable` на не-конфиг классах — критическое нарушение
+
+**Правило** (стандарт §3.1):
+
+> «Декоратор `@gin.configurable` **только на классах** [конфигурации, *не на бизнес-логике*]»
+
+**Что нарушает:**
+
+| Файл | Класс | Параметры (которые проникают в `.gin`) |
+|---|---|---|
+| `src/losses/multi_infonce.py` | `InfoNCELoss` | `temperature_init`, `temperature_final`, `label_smoothing`, `weight_q2g`, `weight_g2q`, `learnable_temperature`, `tau_min`, `tau_max`, `hard_mining_k` |
+| `src/losses/weighted_infonce.py` | `WeightedInfoNCELoss` | `temperature_init`, `learnable_temperature`, `label_smoothing`, `k`, `tau_min`, `tau_max` |
+| `src/datasets/visloc_with_captions.py` (v2) | `GeoLocCaptionDataset` | `query_file`, `data_root`, `image_transform`, `drop_caption_prob`, `seed` |
+
+**Почему это нарушение:** конфиг и бизнес-логика — разные слои. Когда `nn.Module` или `Dataset` декорированы `@gin.configurable`, гин лезет в **их** `__init__` помимо лезения в `TrainConfigGTAUAV.__init__`. В `gtauav_balanced.gin` это видно прямо:
+
+```gin
+# Параметры дублируются — раз в TrainConfigGTAUAV, раз в InfoNCELoss:
+TrainConfigGTAUAV.tau_init = 0.07
+TrainConfigGTAUAV.label_smoothing = 0.1
+TrainConfigGTAUAV.weight_q2g = 0.6
+TrainConfigGTAUAV.weight_g2q = 0.4
+
+# ---- InfoNCELoss (gin-configurable) ----
+InfoNCELoss.temperature_init = 0.07          # ← дубль tau_init
+InfoNCELoss.label_smoothing = 0.1            # ← дубль
+InfoNCELoss.weight_q2g = 0.6                 # ← дубль
+InfoNCELoss.weight_g2q = 0.4                 # ← дубль
+InfoNCELoss.tau_min = 0.01
+InfoNCELoss.tau_max = 0.1
+InfoNCELoss.hard_mining_k = 0
+```
+
+**Это активный источник тихих багов:** если кто-то поменяет `TrainConfigGTAUAV.tau_init = 0.05`, а `InfoNCELoss.temperature_init` забудет — обучение пойдёт с `0.05` в логике trainer-а и `0.07` в самой loss-функции. Никаких ошибок не будет, метрики просто будут странными.
+
+**Масштаб:** **3 класса**, но самый болезненный — `InfoNCELoss`, потому что у него **самая большая зона перекрытия** с `TrainConfigGTAUAV`.
+
+---
+
+### 3. `argparse` — критическое нарушение
+
+**Правило** (стандарт §3.4):
+
+> «**Запрещён argparse** — все параметры из .gin файлов»
+
+**Что нарушает:**
+
+| Файл | Кол-во CLI флагов | Самые проблемные |
+|---|---|---|
+| `src/training/train_gtauav.py::main` | ~15 | `--config`, `--baseline`, `--batch-size`, `--epochs`, `--filter-meta`, `--wandb`, `--gradcam`, `--profile`, `--gin-param`, `--resume`, `--output-dir` |
+| `src/training/train.py::main` (v2) | 1 | `--config` |
+| `scripts/make_split.py::main` | 3 | `--ratio`, `--seed`, `--output-dir` |
+| `scripts/filter_segmentation.py::main` | 3 | `--segm-root`, `--threshold`, `--output` |
+
+**Самый ядовитый паттерн** в `train_gtauav.py`:
+
+```python
+parser.add_argument("--gin-param", nargs="*", help="Override gin params from CLI")
+# ...
+gin.parse_config_files_and_bindings([cfg_file], extra_bindings)
+```
+
+CLI **перекрывает** gin-биндинги. Это создаёт **3 источника правды** на один и тот же параметр: дефолт в `__init__`, значение в `.gin`, значение в `--gin-param`. Какое из них применилось в конкретном запуске — невозможно установить иначе как чтением логов.
+
+**Масштаб:** **тренировочный код** (`train_gtauav.py`) — критично; **скрипты** (`make_split.py`, `filter_segmentation.py`) — пограничный случай (см. ниже §7).
+
+---
+
+### 4. `include` для композиции `.gin` — пограничный случай
+
+**Правило** (`REQUIREMENTS_GIN_STYLE.md` §8):
+
+> «Использовать **минимальный** набор возможностей gin: `@gin.configurable`, `gin.parse_config_file()`. Нет: `gin.register()`, `gin.constant()`, `gin.query_parameter()`, **макросы, ссылки между конфигами**»
+
+**Что есть сейчас:**
+
+```gin
+# conf/gtauav_balanced_asym.gin
+include 'conf/gtauav_balanced.gin'
+
+TrainConfigGTAUAV.shared_encoder = False
+TrainConfigGTAUAV.mona_last_n_blocks = 24
+TrainConfigGTAUAV.output_dir = "out/gtauav/balanced_asym"
+```
+
+`include` — это **встроенная фича gin** для композиции. Формально она:
+- ✅ Не входит в явный список запрещённого (`gin.constant`, `gin.register`, `gin.query_parameter`, макросы)
+- ❌ Попадает под формулировку «**ссылки между конфигами**» в `REQUIREMENTS_GIN_STYLE.md` §8
+
+**По букве правил это нарушение.** По духу — `include` решает реальную проблему (DRY), но создаёт неявную зависимость: чтобы понять, с какими параметрами запускается `gtauav_balanced_stripnet_unfrozen.gin`, нужно прочитать **3 файла** (`unfrozen` → `stripnet` → `balanced`).
+
+**Масштаб:** все 14 `gtauav_*.gin` (кроме `gtauav_balanced.gin`, который сам — корень дерева) и v2 (`baseline_no_text.gin`, `text_heavy.gin` тянут `balanced.gin`).
+
+---
+
+### 5. Один мега-конфиг на всё — не «жёсткое» нарушение, но против духа
+
+**Правило** (стандарт §3.3, `Рекомендуемые_gin-config_категории.md`):
+
+> «Каждый `.gin` → один конфиг-класс»
+>
+> «Принцип разделения: если два параметра меняются **вместе** — в одном конфиге. Если **независимо** — в разных»
+
+**Что есть сейчас:**
+
+`TrainConfigGTAUAV` миксует **несколько независимых осей изменчивости** в один класс:
+- Пути к данным/моделям (меняются при смене машины) ↔ training schedule (меняются при смене эксперимента) ↔ model arch (`baseline_mode`, `shared_encoder`, `mona_*`, `stripnet_*`) ↔ loss (`tau_init`, `label_smoothing`, `weight_q2g`, ...) ↔ sampler (`sampler_type`, `dss_*`) ↔ tracking (`use_wandb`, `use_tb`, `use_gradcam`, `use_profiler`)
+
+В одном `.gin` лежат биндинги для всех этих осей. Когда нужно «то же обучение, но без wandb» — приходится копировать целый `.gin` файл.
+
+**Это не нарушение явного запрета, но прямое следствие нарушения #1:** пока есть один большой `@dataclass + @gin.configurable`, иначе расположить параметры просто негде.
+
+**Масштаб:** проникает во все 12 `gtauav_*.gin` файлов одинаково.
+
+---
+
+### 6. Module-level хардкод путей — нарушение
+
+**Правило** (стандарт §6 чеклист):
+
+> «Нет захардкоженных model ID / промптов / размеров?»
+
+**Что нарушает:**
+
+```python
+# src/training/train_gtauav.py (module level)
+_RGB_ROOT = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR"
+_CAPTION_ROOT = "/media/servml/SSD_2_2TB/datasets/cvgl_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"
+
+# src/datasets/gtauav_dataset.py (module level) — ДУБЛЬ:
+_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
+_CAPTION_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions")
+
+# scripts/make_split.py (module level):
+_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
+
+# scripts/filter_segmentation.py (module level):
+SEGM_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm")
+```
+
+**Главная проблема — не сам факт хардкода**, а то, что один и тот же путь `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR` дублируется в **трёх** местах: train_gtauav.py + gtauav_dataset.py + make_split.py. Если переехать на другую машину — нужно править 3 файла; если забыть один — silent breakage.
+
+**Масштаб:** 4 файла, ~10 module-level констант пути.
+
+---
+
+### 7. Скрипты `make_split.py` / `filter_segmentation.py` — спорный случай
+
+**Что есть:**
+- `argparse` (3 параметра в каждом)
+- module-level пути (`_RGB_ROOT`, `SEGM_ROOT`)
+- Запускаются однократно перед тренировкой
+- Не интегрированы в `Trainer.run()`
+
+**Применять ли стандарт «нет argparse» к ним?**
+
+Стандарт направлен против `argparse` в **тренировочном коде**, где параметров 50+ и нужна воспроизводимость. У `make_split.py` 3 параметра, и это **разовая препроцессинг-утилита**. Здесь буква и дух стандарта расходятся.
+
+**Это не «несоответствие, требующее немедленной правки»**, а **открытый вопрос для следующего шага**: применять ли gin-стиль к препроцессинг-скриптам или оставить argparse.
+
+---
+
+### 8. Sofia models — `@dataclass` без gin (формально не нарушение)
+
+**Что есть:**
+
+```python
+# src/models/sofia_v1/config.py
+@dataclass
+class SOFIAv1Config:
+    variant: Literal["tiny_tiny", "tiny", "small", "small_v2"] = "small"
+    in_channels: int = 3
+    input_size: int = 256
+    dcn_variant: Literal["v2", "v4"] = "v2"
+    d_descriptor: int = 1024
+    use_film_altitude: bool = True
+    altitude_norm: float = 500.0
+    use_text_film_uav: bool = True
+    use_text_film_sat: bool = True
+    text_film_dim: int = 1024
+    text_film_hidden: int = 256
+```
+
+**Стандарт** запрещает **`@gin.configurable + @dataclass`**, но не запрещает `@dataclass` сам по себе. `SOFIAv1Config` без gin — формально стандарт **не нарушает**.
+
+**Однако:** в `gtauav_balanced_sofia*.gin` (по скриншоту локально есть, в репо пока нет) параметры sofia, очевидно, прокидываются в `TrainConfigGTAUAV` как обычные поля. То есть sofia сейчас живёт в **двух** местах:
+- `SOFIAv1Config` (dataclass, в коде модели)
+- `TrainConfigGTAUAV.sofia_*` (если там такие есть) или только через позиционное создание `SOFIAv1Config()` в `Trainer`
+
+Это не нарушение, но **архитектурный раскол**: dataclass-конфиги внутри модельных подсистем + gin-конфиги снаружи.
+
+**Это тоже открытый вопрос для следующего шага**, не текущая проблема.
+
+---
+
+### 9. Прочее по чеклисту стандарта (мелочи)
+
+| Пункт | Состояние | Файлы |
+|---|---|---|
+| `from __future__ import annotations` первой строкой | ✅ есть в ключевых файлах | проверено в `train_gtauav.py`, `gtauav_dataset.py`, `multi_infonce.py`, `weighted_infonce.py`, `train.py` |
+| Строгие type hints | ✅ в основном | пара мест с `dict` без параметров (`_atomic_save(obj: dict)`) |
+| Google-style docstrings | ✅ есть, качество хорошее | — |
+| `@torch.inference_mode()` вместо `@torch.no_grad()` | ❌ используется `@torch.no_grad()` | `train_gtauav.py::_evaluate`, `_embed_drone_queries` |
+| Atomic writes | ⚠️ есть, но без cleanup на ошибке | `train_gtauav.py::_atomic_save` (нет `try/except` — `.tmp` остаётся при сбое) |
+| Английский в коде/комментариях | ✅ есть | — |
+| Импорты stdlib → third-party → local | ✅ есть | — |
+
+---
+
+## Сводка — что нарушено и насколько срочно
+
+| # | Нарушение | Срочность | Зона воздействия |
+|---|---|---|---|
+| 1 | `@gin.configurable + @dataclass` на `TrainConfigGTAUAV` | 🔴 критично | блокирует всё остальное |
+| 2 | `@gin.configurable` на `InfoNCELoss`, `WeightedInfoNCELoss`, `GeoLocCaptionDataset` | 🔴 критично | активный источник тихих багов |
+| 3 | `argparse` в тренировочном коде | 🔴 критично | три источника правды на параметр |
+| 4 | `include` между `.gin` файлами | 🟡 пограничный | формально нарушает «нет ссылок между конфигами» |
+| 5 | Один мега-конфиг (нет разделения на оси) | 🟡 следствие #1 | разрешится с #1 |
+| 6 | Module-level пути в 4 файлах | 🟡 нарушение, но не критично | дубли — реальная проблема |
+| 7 | `argparse` в скриптах препроцессинга | ⚪ открытый вопрос | спорный случай |
+| 8 | `@dataclass` в `SOFIAv1Config` (без gin) | ⚪ открытый вопрос | формально не нарушение |
+| 9 | `@torch.no_grad()` вместо `@torch.inference_mode()` | 🟢 мелочь | косметика |
+| 10 | `_atomic_save` без cleanup .tmp | 🟢 мелочь | редкое последствие |
+
+---
+
+## Что предлагаю обсудить дальше (Шаг 2)
+
+Прежде чем двигать код, нужно принять решения по 4 развилкам:
+
+1. **`include` между `.gin`** — терпим как удобный DRY-механизм или приводим к плоской иерархии (каждый эксперимент = самодостаточный набор `.gin` без `include`)?
+
+2. **Разделение `TrainConfigGTAUAV` на классы** — на сколько и по каким осям?
+   - Вариант **5 классов** (Pipeline / Hardware / Models / Training / Tracking) из «`Рекомендуемые_gin-config_категории.md`».
+   - Вариант **6 классов** (отдельный Loss + Sampler).
+   - Вариант **3 класса** (всё, что было — в Training, плюс Pipeline + Tracking).
+
+3. **Sofia + dataclass** — оставлять `SOFIAv1Config` как dataclass-структуру внутри модели или переписать в обычный класс/влить в `ModelsConfig`?
+
+4. **Скрипты** — переводить на gin или оставить argparse?
+
+После решения этих 4 пунктов план рефакторинга становится однозначным.

scripts/filter_segmentation.py

@@ -25,7 +25,7 @@ from tqdm import tqdm
 
 LOGGER = logging.getLogger("caption_test.filter_seg")
 
-SEGM_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-aug/segm")
+SEGM_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm")
 EXCLUDE_CLASSES = {0, 4}  # background, water
 DEFAULT_THRESHOLD = 0.90
 

scripts/make_split.py

@@ -19,7 +19,7 @@ import coloredlogs
 
 LOGGER = logging.getLogger("caption_test.make_split")
 
-_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
+_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
 
 
 def main() -> None:

src/conf/__init__.py

@@ -0,0 +1,54 @@
+"""Gin-configurable settings for the caption-test project.
+
+Five universal axes of variability:
+  - PipelineConfig     — paths, training schedule, output, resume
+  - HardwareConfig     — batch size, accumulation, AMP, gradient checkpointing
+  - TrainingConfig     — loss + optimizer + sampler (the recipe)
+  - TrackingConfig     — wandb / tensorboard / gradcam / profiler
+
+Plus a model-family config: ModelsCommonConfig describes the active backbone,
+and one of {DINOv3, StripNet, SOFIAv1, SOFIAv71} ModelsConfig classes
+parameterises it.
+
+Plus PreprocessConfig used only by scripts/make_split.py and
+scripts/filter_segmentation.py.
+
+All configs are loaded together via load_all_configs(path2cfg) — see
+config_loader.py.
+"""
+
+from src.conf.config_loader import load_all_configs
+from src.conf.hardware_conf import HardwareConfig, get_hardware_cfg
+from src.conf.models_common_conf import ModelsCommonConfig, get_models_common_cfg
+from src.conf.models_dinov3_conf import DINOv3ModelsConfig, get_models_dinov3_cfg
+from src.conf.models_sofia_v1_conf import SOFIAv1ModelsConfig, get_models_sofia_v1_cfg
+from src.conf.models_sofia_v71_conf import SOFIAv71ModelsConfig, get_models_sofia_v71_cfg
+from src.conf.models_stripnet_conf import StripNetModelsConfig, get_models_stripnet_cfg
+from src.conf.pipeline_conf import PipelineConfig, get_pipeline_cfg
+from src.conf.preprocess_conf import PreprocessConfig, get_preprocess_cfg
+from src.conf.tracking_conf import TrackingConfig, get_tracking_cfg
+from src.conf.training_conf import TrainingConfig, get_training_cfg
+
+__all__ = [
+    "DINOv3ModelsConfig",
+    "HardwareConfig",
+    "ModelsCommonConfig",
+    "PipelineConfig",
+    "PreprocessConfig",
+    "SOFIAv1ModelsConfig",
+    "SOFIAv71ModelsConfig",
+    "StripNetModelsConfig",
+    "TrackingConfig",
+    "TrainingConfig",
+    "get_hardware_cfg",
+    "get_models_common_cfg",
+    "get_models_dinov3_cfg",
+    "get_models_sofia_v1_cfg",
+    "get_models_sofia_v71_cfg",
+    "get_models_stripnet_cfg",
+    "get_pipeline_cfg",
+    "get_preprocess_cfg",
+    "get_tracking_cfg",
+    "get_training_cfg",
+    "load_all_configs",
+]

src/conf/config_loader.py

@@ -0,0 +1,171 @@
+"""Single entry point for loading all configs in a training run.
+
+Layout (per REQUIREMENTS_GIN_STYLE.md §1, extended for multi-experiment):
+
+    in/config_files/
+        ├── training.gin                    # common: shared by all training presets
+        ├── hardware_default.gin            # common: for DINOv3 / StripNet
+        ├── hardware_no_gradckpt.gin        # common: for SOFIA backbones
+        ├── tracking.gin                    # common: shared by all
+        └── <preset_name>/
+            ├── pipeline.gin                # local: output_dir, paths, schedule
+            └── models.gin                  # local: backbone-specific bindings
+
+`load_all_configs` parses common files at the path2cfg root + preset locals
+in `path2cfg/<preset_name>/`. Two-pass loading:
+  Pass 1: read <preset_name>/models.gin to learn the backbone.
+  Pass 2: parse common files + preset locals in one batch.
+
+Override semantics: if the preset directory contains its own training.gin /
+hardware.gin / tracking.gin, those files are appended AFTER the common
+versions, so their bindings win (gin: last write wins).
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import gin
+
+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
+
+logger = logging.getLogger(__name__)
+
+# Maps ModelsCommonConfig.backbone → family-specific config class.
+_BACKBONE_TO_MODELS_CLS = {
+    "dinov3":    DINOv3ModelsConfig,
+    "stripnet":  StripNetModelsConfig,
+    "sofia_v1":  SOFIAv1ModelsConfig,
+    "sofia_v71": SOFIAv71ModelsConfig,
+}
+
+# Sofia backbones disable gradient checkpointing.
+_NO_GRADCKPT_BACKBONES = {"sofia_v1", "sofia_v71"}
+
+# Common filenames at the path2cfg root.
+_COMMON_TRAINING = "training.gin"
+_COMMON_TRACKING = "tracking.gin"
+_COMMON_HARDWARE_DEFAULT = "hardware_default.gin"
+_COMMON_HARDWARE_NO_GRADCKPT = "hardware_no_gradckpt.gin"
+
+# Files a preset must always have locally (in <path2cfg>/<preset_name>/).
+_LOCAL_REQUIRED = ("pipeline.gin", "models.gin")
+
+# Files that can optionally be overridden locally; if present in the preset
+# directory, they win over the common version.
+_LOCAL_OVERRIDABLE = ("training.gin", "hardware.gin", "tracking.gin")
+
+
+def load_all_configs(path2cfg: str, preset_name: str) -> dict[str, Any]:
+    """Parse common gin files + preset gin files and return all configs.
+
+    Args:
+        path2cfg: Path to in/config_files/ (per REQUIREMENTS_GIN_STYLE.md §5).
+                  Must contain training.gin, hardware_default.gin,
+                  hardware_no_gradckpt.gin, tracking.gin at its root, plus
+                  per-preset subdirectories.
+        preset_name: Name of the preset subdirectory under path2cfg, e.g.
+                     'gtauav_balanced'.
+
+    Returns:
+        Dict with keys: 'pipeline', 'hardware', 'training', 'tracking',
+        'models_common', 'models'.
+
+    Raises:
+        FileNotFoundError: If path2cfg, preset directory, or a required file is missing.
+        ValueError: If models_common.backbone is not a known value.
+    """
+    common_dir = Path(path2cfg)
+    preset_dir = common_dir / preset_name
+
+    # Sanity checks.
+    if not common_dir.is_dir():
+        raise FileNotFoundError(
+            f"Config root not found: {common_dir}. "
+            f"Per REQUIREMENTS_GIN_STYLE.md §5, this should be "
+            f"<proj_dir>/in/config_files/.",
+        )
+    if not preset_dir.is_dir():
+        raise FileNotFoundError(
+            f"Preset directory not found: {preset_dir}. "
+            f"Available presets in {common_dir}: "
+            f"{sorted(d.name for d in common_dir.iterdir() if d.is_dir())}",
+        )
+    for required in _LOCAL_REQUIRED:
+        if not (preset_dir / required).is_file():
+            raise FileNotFoundError(
+                f"Preset {preset_name} is missing required file '{required}' "
+                f"(looked in {preset_dir})",
+            )
+
+    # ===== Pass 1: peek at models.gin to learn the backbone. =====
+    gin.clear_config()
+    gin.parse_config_file(str(preset_dir / "models.gin"))
+    backbone = ModelsCommonConfig().backbone
+    if backbone not in _BACKBONE_TO_MODELS_CLS:
+        raise ValueError(
+            f"Unknown backbone={backbone!r} in {preset_dir / 'models.gin'}; "
+            f"expected one of {sorted(_BACKBONE_TO_MODELS_CLS)}",
+        )
+
+    # ===== Pass 2: build the full file list and parse in one batch. =====
+    # Order: common first, preset locals last (so locals win on overrides).
+    gin_files: list[Path] = [
+        common_dir / _COMMON_TRAINING,
+        common_dir / (
+            _COMMON_HARDWARE_NO_GRADCKPT if backbone in _NO_GRADCKPT_BACKBONES
+            else _COMMON_HARDWARE_DEFAULT
+        ),
+        common_dir / _COMMON_TRACKING,
+        # Always-required preset locals.
+        preset_dir / "pipeline.gin",
+        preset_dir / "models.gin",
+    ]
+
+    # Optional preset overrides (rare).
+    for overridable in _LOCAL_OVERRIDABLE:
+        local = preset_dir / overridable
+        if local.is_file():
+            gin_files.append(local)
+            logger.info("Preset %s overrides %s locally", preset_name, overridable)
+
+    # Sanity: all chosen files must exist.
+    for f in gin_files:
+        if not f.is_file():
+            raise FileNotFoundError(f"Required gin file not found: {f}")
+
+    # MANDATORY: clear gin global state before parsing.
+    gin.clear_config()
+    gin.parse_config_files_and_bindings(
+        config_files=[str(f) for f in gin_files],
+        bindings=[],
+    )
+    logger.info(
+        "Loaded preset %s with %d gin files (backbone=%s)",
+        preset_name, len(gin_files), backbone,
+    )
+
+    # Build all configs from gin global state.
+    models_common = ModelsCommonConfig()
+    models_specific = _BACKBONE_TO_MODELS_CLS[models_common.backbone]()
+
+    return {
+        "pipeline":      PipelineConfig(),
+        "hardware":      HardwareConfig(),
+        "training":      TrainingConfig(),
+        "tracking":      TrackingConfig(),
+        "models_common": models_common,
+        "models":        models_specific,
+    }
+   
+

src/conf/hardware_conf.py

@@ -0,0 +1,46 @@
+"""GPU profile + memory/compute optimisation flags."""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class HardwareConfig:
+    """Hardware-bound parameters: VRAM footprint and throughput.
+
+    These do not change the training recipe (loss/optimizer/sampler), only
+    how many samples fit on the device.
+    """
+
+    def __init__(
+        self,
+        device: str = "cuda",
+        batch_size: int = 8,
+        grad_accum_steps: int = 8,
+        num_workers: int = 4,
+        use_amp: bool = True,
+        gradient_checkpointing: bool = True,
+        reserve_gb: float = 2.0,
+        total_vram_gb: float = 24.0,
+    ) -> None:
+        self.device = device
+        self.batch_size = batch_size
+        self.grad_accum_steps = grad_accum_steps
+        self.num_workers = num_workers
+        self.use_amp = use_amp
+        self.gradient_checkpointing = gradient_checkpointing
+        self.reserve_gb = reserve_gb
+        self.total_vram_gb = total_vram_gb
+        # Derived.
+        self.available_vram_gb = self.total_vram_gb - self.reserve_gb
+        self.effective_batch_size = self.batch_size * self.grad_accum_steps
+
+
+def get_hardware_cfg(path2cfg: str) -> HardwareConfig:
+    """Load ONLY hardware config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}hardware.gin")
+    return HardwareConfig()
+
+

src/conf/models_common_conf.py

@@ -0,0 +1,50 @@
+"""Backbone-agnostic model parameters.
+
+`backbone` selects which family-specific Models config is loaded by
+config_loader.load_all_configs.
+"""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class ModelsCommonConfig:
+    """Shared model fields across all backbones.
+
+    `backbone` is the dispatch key — one of:
+      - 'dinov3'    → DINOv3ModelsConfig
+      - 'stripnet'  → StripNetModelsConfig
+      - 'sofia_v1'  → SOFIAv1ModelsConfig
+      - 'sofia_v71' → SOFIAv71ModelsConfig
+
+    `baseline_mode=True` disables text fusion entirely (gates locked at 1.0,
+    DGTRS-CLIP not loaded, TextFusionMLP not built). Used for Δ R@1 baselines.
+
+    `init_gate` controls the initial sigmoid value of GatedFusion gates
+    (0.7 = 70% image, 30% text by default; 0.3 = text-heavy; 0.9 = image-heavy).
+
+    `lrsclip_path` is the path to the DGTRS-CLIP checkpoint (only loaded when
+    text fusion is active).
+    """
+
+    def __init__(
+        self,
+        backbone: str = "dinov3",
+        baseline_mode: bool = False,
+        init_gate: float = 0.7,
+        lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
+    ) -> None:
+        self.backbone = backbone
+        self.baseline_mode = baseline_mode
+        self.init_gate = init_gate
+        self.lrsclip_path = lrsclip_path
+
+
+def get_models_common_cfg(path2cfg: str) -> ModelsCommonConfig:
+    """Load ONLY models_common config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}models.gin")
+    return ModelsCommonConfig()
+

src/conf/models_dinov3_conf.py

@@ -0,0 +1,44 @@
+"""DINOv3 backbone configuration: encoders + MONA adapters."""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class DINOv3ModelsConfig:
+    """DINOv3 ViT-L/16 with MONA adapters (CVPR 2025).
+
+    `shared_encoder=True` uses a single DINOv3 WEB encoder for both drone and
+    satellite branches (default; ~432M params total). When False, separate WEB
+    (drone) + SAT (satellite) encoders are built (~733M params total, +4-5GB
+    VRAM).
+
+    MONA adapters are injected in the LAST `mona_last_n_blocks` of the 24
+    ViT blocks (default: 12 = top half). Set to 24 for full-capacity asymmetric
+    setup.
+    """
+
+    def __init__(
+        self,
+        dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
+        dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
+        shared_encoder: bool = True,
+        mona_bottleneck: int = 64,
+        mona_last_n_blocks: int = 12,
+        lora_rank: int = 4,
+    ) -> None:
+        self.dino_web_path = dino_web_path
+        self.dino_sat_path = dino_sat_path
+        self.shared_encoder = shared_encoder
+        self.mona_bottleneck = mona_bottleneck
+        self.mona_last_n_blocks = mona_last_n_blocks
+        self.lora_rank = lora_rank
+
+
+def get_models_dinov3_cfg(path2cfg: str) -> DINOv3ModelsConfig:
+    """Load ONLY DINOv3 models config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}models.gin")
+    return DINOv3ModelsConfig()
+

src/conf/models_sofia_v1_conf.py

@@ -0,0 +1,69 @@
+"""SOFIA v1 backbone configuration: 4-stage StripNet+DCNv4 + GGeM heads."""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class SOFIAv1ModelsConfig:
+    """SOFIA v1: lightweight StripNet+DCNv4 backbone + heads.
+
+    `variant_label` chooses backbone size (architecture dimensions are
+    resolved inside the model code from this label):
+        tiny_tiny: dims [16, 32, 80, 128]   (~0.4M)
+        tiny     : dims [32, 64, 128, 256]  (~1M)
+        small    : dims [64, 128, 320, 512] (~3M, default in code)
+        small_v2 : dims [64, 128, 256, 384] (~2M)
+
+    `dcn_variant`: 'v2' = torchvision DeformConv2d (stable). 'v4' = OpenGVLab
+    DCNv4 (faster but ~9 MB / forward leak from C++ extension).
+
+    Text fusion is two-level:
+      - Mid-level: Text-FiLM modulates feature maps before GGeM (when
+        use_text_film_uav / use_text_film_sat = True).
+      - Late-level: GatedFusion on descriptors (handled outside this config).
+    """
+
+    def __init__(
+        self,
+        # ---- Backbone ----
+        variant_label: str = "small",      # 'tiny_tiny' | 'tiny' | 'small' | 'small_v2'
+        in_channels: int = 3,
+        input_size: int = 256,
+        dcn_variant: str = "v2",           # 'v2' | 'v4'
+        # ---- Heads ----
+        d_descriptor: int = 1024,
+        return_normalized: bool = False,
+        # ---- Altitude-FiLM (UAV head) ----
+        use_film_altitude: bool = True,
+        altitude_norm: float = 500.0,
+        # ---- Text-FiLM (mid-level fusion) ----
+        use_text_film_uav: bool = True,
+        use_text_film_sat: bool = True,
+        text_film_dim: int = 1024,
+        text_film_hidden: int = 256,
+        # ---- LoRA on DGTRS-CLIP text encoder ----
+        lora_rank: int = 4,
+    ) -> None:
+        self.variant_label = variant_label
+        self.in_channels = in_channels
+        self.input_size = input_size
+        self.dcn_variant = dcn_variant
+        self.d_descriptor = d_descriptor
+        self.return_normalized = return_normalized
+        self.use_film_altitude = use_film_altitude
+        self.altitude_norm = altitude_norm
+        self.use_text_film_uav = use_text_film_uav
+        self.use_text_film_sat = use_text_film_sat
+        self.text_film_dim = text_film_dim
+        self.text_film_hidden = text_film_hidden
+        self.lora_rank = lora_rank
+
+
+def get_models_sofia_v1_cfg(path2cfg: str) -> SOFIAv1ModelsConfig:
+    """Load ONLY SOFIA v1 models config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}models.gin")
+    return SOFIAv1ModelsConfig()
+

src/conf/models_sofia_v71_conf.py

@@ -0,0 +1,178 @@
+"""SOFIA v7.1 backbone: 4-stage StripDCN + MambaVision + CVGL-Aware Head."""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class SOFIAv71ModelsConfig:
+    """SOFIA v7.1 student model.
+
+    Mirrors src/models/sofia_v71/config.py::SOFIAConfig with one
+    difference: `mamba_extra_kwargs` (a dict in the dataclass) is flattened
+    into 5 explicit fields here, and reassembled into a dict for downstream
+    code.
+
+    Variant scale presets (see model code):
+        Tiny: stem=16/32, dims=[48, 96, 176, 224],  depths=[2, 3, 4, 2]   (~5M)
+        M   : stem=64/128, dims=[256, 512, 1280, 1536], depths=[3, 4, 15, 3] (~500M, default)
+        L   : stem=64/128, dims=[256, 512, 1536, 2048], depths=[3, 4, 20, 3] (~1B)
+
+    For the active experiment (Tiny preset, see `presets/gtauav_balanced_sofia/`)
+    you can override individual fields directly without resorting to a
+    'preset' string parameter — every architectural dimension is bindable.
+
+    Tiny needs `num_heads_*=4` (channels 176/224 not divisible by 8) and
+    `mamba_headdim=16` (channels not divisible by 64).
+    """
+
+    def __init__(
+        self,
+        # ---- Variant label (informational, used in logs/checkpoints) ----
+        variant_label: str = "M",                # 'M' | 'L' | 'Tiny'
+        # ---- Input ----
+        input_size: int = 256,
+        in_channels: int = 3,
+        # ---- Stem ----
+        stem_mid: int = 64,
+        stem_out: int = 128,
+        # ---- Backbone dimensions (per stage s1..s4) ----
+        # Lists default to None; concrete defaults are filled in __init__ to
+        # avoid the def f(x=[]) anti-pattern.
+        embed_dims: list[int] | None = None,     # default [256, 512, 1280, 1536] (M)
+        depths: list[int] | None = None,         # default [3, 4, 15, 3] (M)
+        # ---- Stage 1-2 block params ----
+        mbconv_expand: int = 4,
+        se_ratio: int = 16,
+        strip_kernel_s1: int = 7,
+        strip_kernel_s2: int = 5,
+        mix_kernels: list[int] | None = None,    # default [3, 5, 7]
+        use_dcn_strip: bool = True,
+        # ---- Stage 3-4 (MambaVision) ----
+        mamba_d_state: int = 16,
+        mamba_dt_rank: int | None = None,        # auto = max(1, C // 16)
+        mamba_backend: str = "auto",             # 'auto' | 'torch' | 'mamba_ssm'
+        mamba_variant: str = "mamba2",           # 'mamba1' | 'mamba2' | 'efficient_vmamba'
+        # mamba_extra_kwargs flattened (assembled back into a dict in __init__):
+        mamba_d_state_mamba2: int = 64,
+        mamba_headdim: int = 64,
+        mamba_expand: int = 2,
+        mamba_d_conv: int = 4,
+        mamba_n_directions: int = 2,
+        # ---- Heads / attention ----
+        num_heads_s3: int = 8,
+        num_heads_s4: int = 8,
+        use_strip_branch_s3: bool = True,
+        use_strip_branch_s4: bool = False,
+        ffn_expand: int = 4,
+        # ---- EVSS bridge ----
+        use_evss_bridge: bool = False,
+        evss_bridge_locations: list[str] | None = None,   # default ['pre_stage3']
+        # ---- Neck ----
+        neck_channels: int = 192,
+        # ---- CVGL-Aware Head v7.1-α ----
+        d_descriptor: int = 512,
+        use_asymmetric_heads: bool = True,
+        chp_rings: int = 8,
+        chp_angles: int = 16,
+        chp_harmonics: int = 4,
+        use_film_altitude: bool = True,
+        altitude_norm: float = 500.0,
+        ring_count: int = 4,
+        use_ring_aux: bool = True,
+        # ---- Text fusion ----
+        return_normalized: bool = True,
+        use_text_film_sat: bool = False,
+        use_text_film_uav: bool = False,
+        text_film_dim: int = 1024,
+        text_film_hidden: int = 256,
+        # ---- Weight-sharing ----
+        share_stages_1_2: bool = True,
+        # ---- KD taps ----
+        enable_kd_taps: bool = True,
+        # ---- Deployment ----
+        precision: str = "fp16",                 # 'fp32' | 'fp16' | 'int8_mixed'
+        # ---- LoRA on DGTRS-CLIP text encoder ----
+        lora_rank: int = 4,
+    ) -> None:
+        # Variant label.
+        self.variant_label = variant_label
+        # Input.
+        self.input_size = input_size
+        self.in_channels = in_channels
+        # Stem.
+        self.stem_mid = stem_mid
+        self.stem_out = stem_out
+        # Backbone dimensions.
+        self.embed_dims = embed_dims if embed_dims is not None else [256, 512, 1280, 1536]
+        self.depths = depths if depths is not None else [3, 4, 15, 3]
+        # Stage 1-2.
+        self.mbconv_expand = mbconv_expand
+        self.se_ratio = se_ratio
+        self.strip_kernel_s1 = strip_kernel_s1
+        self.strip_kernel_s2 = strip_kernel_s2
+        self.mix_kernels = mix_kernels if mix_kernels is not None else [3, 5, 7]
+        self.use_dcn_strip = use_dcn_strip
+        # Stage 3-4.
+        self.mamba_d_state = mamba_d_state
+        self.mamba_dt_rank = mamba_dt_rank
+        self.mamba_backend = mamba_backend
+        self.mamba_variant = mamba_variant
+        self.mamba_d_state_mamba2 = mamba_d_state_mamba2
+        self.mamba_headdim = mamba_headdim
+        self.mamba_expand = mamba_expand
+        self.mamba_d_conv = mamba_d_conv
+        self.mamba_n_directions = mamba_n_directions
+        # Heads.
+        self.num_heads_s3 = num_heads_s3
+        self.num_heads_s4 = num_heads_s4
+        self.use_strip_branch_s3 = use_strip_branch_s3
+        self.use_strip_branch_s4 = use_strip_branch_s4
+        self.ffn_expand = ffn_expand
+        # EVSS.
+        self.use_evss_bridge = use_evss_bridge
+        self.evss_bridge_locations = (
+            evss_bridge_locations if evss_bridge_locations is not None else ["pre_stage3"]
+        )
+        # Neck.
+        self.neck_channels = neck_channels
+        # CVGL Head.
+        self.d_descriptor = d_descriptor
+        self.use_asymmetric_heads = use_asymmetric_heads
+        self.chp_rings = chp_rings
+        self.chp_angles = chp_angles
+        self.chp_harmonics = chp_harmonics
+        self.use_film_altitude = use_film_altitude
+        self.altitude_norm = altitude_norm
+        self.ring_count = ring_count
+        self.use_ring_aux = use_ring_aux
+        # Text fusion.
+        self.return_normalized = return_normalized
+        self.use_text_film_sat = use_text_film_sat
+        self.use_text_film_uav = use_text_film_uav
+        self.text_film_dim = text_film_dim
+        self.text_film_hidden = text_film_hidden
+        # Sharing / KD / deploy.
+        self.share_stages_1_2 = share_stages_1_2
+        self.enable_kd_taps = enable_kd_taps
+        self.precision = precision
+        # LoRA.
+        self.lora_rank = lora_rank
+        # Derived: assemble mamba_extra_kwargs back for downstream consumers.
+        self.mamba_extra_kwargs = {
+            "d_state_mamba2": self.mamba_d_state_mamba2,
+            "headdim": self.mamba_headdim,
+            "expand": self.mamba_expand,
+            "d_conv": self.mamba_d_conv,
+            "n_directions": self.mamba_n_directions,
+        }
+
+
+def get_models_sofia_v71_cfg(path2cfg: str) -> SOFIAv71ModelsConfig:
+    """Load ONLY SOFIA v71 models config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}models.gin")
+    return SOFIAv71ModelsConfig()
+
+

src/conf/models_stripnet_conf.py

@@ -0,0 +1,40 @@
+"""StripNet backbone configuration."""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class StripNetModelsConfig:
+    """StripNet-small encoder with Conv-MONA adaptation.
+
+    `stripnet_freeze=True` keeps the backbone frozen and only trains MONA on
+    the last `stripnet_mona_last_n_stages` of 4 stages.
+
+    `stripnet_freeze=False` (full fine-tune) makes the backbone trainable; in
+    that case backbone params get a separate LR group at
+    `learning_rate * stripnet_backbone_lr_factor` (typically 0.1).
+    """
+
+    def __init__(
+        self,
+        stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
+        stripnet_freeze: bool = True,
+        stripnet_mona_last_n_stages: int = 2,
+        stripnet_backbone_lr_factor: float = 0.1,
+        lora_rank: int = 4,
+    ) -> None:
+        self.stripnet_path = stripnet_path
+        self.stripnet_freeze = stripnet_freeze
+        self.stripnet_mona_last_n_stages = stripnet_mona_last_n_stages
+        self.stripnet_backbone_lr_factor = stripnet_backbone_lr_factor
+        self.lora_rank = lora_rank
+
+
+def get_models_stripnet_cfg(path2cfg: str) -> StripNetModelsConfig:
+    """Load ONLY StripNet models config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}models.gin")
+    return StripNetModelsConfig()
+

src/conf/pipeline_conf.py

@@ -0,0 +1,53 @@
+
+"""Pipeline orchestration: data IO, training schedule, output, resume."""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class PipelineConfig:
+    """What to train on, where to save, and how long.
+
+    All paths are absolute or relative to the project root. Defaults match
+    the servml workstation; override in pipeline.gin for other machines.
+    """
+
+    def __init__(
+        self,
+        # Data inputs.
+        train_json: str = "meta/train_80.json",
+        test_json: str = "meta/test_20.json",
+        rgb_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR",
+        caption_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions",
+        filter_meta: str | None = None,
+        # Training schedule.
+        epochs: int = 10,
+        warmup_epochs: int = 2,
+        eval_every: int = 1,
+        # Reproducibility & output.
+        seed: int = 42,
+        output_dir: str = "out/gtauav/with_text",
+        resume_from: str | None = None,
+    ) -> None:
+        self.train_json = train_json
+        self.test_json = test_json
+        self.rgb_root = rgb_root
+        self.caption_root = caption_root
+        self.filter_meta = filter_meta
+        self.epochs = epochs
+        self.warmup_epochs = warmup_epochs
+        self.eval_every = eval_every
+        self.seed = seed
+        self.output_dir = output_dir
+        self.resume_from = resume_from
+
+
+def get_pipeline_cfg(path2cfg: str) -> PipelineConfig:
+    """Load ONLY pipeline config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}pipeline.gin")
+    return PipelineConfig()
+
+    

src/conf/preprocess_conf.py

@@ -0,0 +1,56 @@
+"""Preprocessing configuration: train/test split + segmentation filter."""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class PreprocessConfig:
+    """Used only by scripts/make_split.py and scripts/filter_segmentation.py.
+
+    Lives in a separate preset (presets/preprocess/preprocess.gin) — it is
+    not consumed by the training pipeline. Held independently from
+    PipelineConfig so that preprocess can run without a training preset
+    being active.
+    """
+
+    def __init__(
+        self,
+        # Inputs.
+        rgb_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR",
+        segm_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm",
+        # make_split.py params.
+        split_ratio: float = 0.8,
+        split_seed: int = 42,
+        split_input_train: str = "cross-area-drone2sate-train.json",
+        split_input_test: str = "cross-area-drone2sate-test.json",
+        split_output_dir: str = "meta",
+        split_output_train: str = "train_80.json",
+        split_output_test: str = "test_20.json",
+        # filter_segmentation.py params.
+        seg_threshold: float = 0.90,
+        seg_exclude_classes: list[int] | None = None,   # default [0, 4]: background + water
+        seg_filter_output: str = "meta/seg_filter.json",
+    ) -> None:
+        self.rgb_root = rgb_root
+        self.segm_root = segm_root
+        self.split_ratio = split_ratio
+        self.split_seed = split_seed
+        self.split_input_train = split_input_train
+        self.split_input_test = split_input_test
+        self.split_output_dir = split_output_dir
+        self.split_output_train = split_output_train
+        self.split_output_test = split_output_test
+        self.seg_threshold = seg_threshold
+        self.seg_exclude_classes = seg_exclude_classes if seg_exclude_classes is not None else [0, 4]
+        self.seg_filter_output = seg_filter_output
+
+
+def get_preprocess_cfg(path2cfg: str) -> PreprocessConfig:
+    """Load preprocess config from the preprocess preset directory."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}preprocess.gin")
+    return PreprocessConfig()
+
+

src/conf/tracking_conf.py

@@ -0,0 +1,51 @@
+"""Experiment tracking + diagnostics.
+
+Independent axis: changing these flags does not affect training results,
+only what is observed/recorded.
+"""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class TrackingConfig:
+    """Wandb / TensorBoard / Grad-CAM / profiler / gradient norms."""
+
+    def __init__(
+        self,
+        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,
+        gradcam_samples: int = 8,
+        use_profiler: bool = False,
+        profiler_warmup: int = 3,
+        profiler_active: int = 5,
+    ) -> None:
+        self.use_wandb = use_wandb
+        self.use_tb = use_tb
+        self.wandb_project = wandb_project
+        self.wandb_run_name = wandb_run_name
+        self.wandb_entity = wandb_entity
+        self.log_grad_norms = log_grad_norms
+        self.use_gradcam = use_gradcam
+        self.gradcam_every = gradcam_every
+        self.gradcam_samples = gradcam_samples
+        self.use_profiler = use_profiler
+        self.profiler_warmup = profiler_warmup
+        self.profiler_active = profiler_active
+
+
+def get_tracking_cfg(path2cfg: str) -> TrackingConfig:
+    """Load ONLY tracking config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}tracking.gin")
+    return TrackingConfig()
+
+

src/conf/training_conf.py

@@ -0,0 +1,93 @@
+"""Training recipe: loss + optimizer + sampler.
+
+Three concerns kept together because they form one coherent recipe — they
+co-vary across experiments. Splitting Loss vs Optimizer vs Sampler can be
+done later if a need emerges.
+"""
+
+from __future__ import annotations
+
+import gin
+
+
+@gin.configurable
+class TrainingConfig:
+    """Loss + optimizer + sampler.
+
+    Selects between InfoNCELoss and WeightedInfoNCELoss via `loss_type`.
+    Selects between DSS / mutex / plain shuffle via `sampler_type`.
+    """
+
+    def __init__(
+        self,
+        # ---- Loss: shared between InfoNCELoss and WeightedInfoNCELoss ----
+        loss_type: str = "symmetric",            # 'symmetric' | 'weighted'
+        tau_init: float = 0.07,
+        tau_min: float = 0.01,
+        tau_max: float = 0.1,
+        learnable_temperature: bool = True,
+        label_smoothing: float = 0.1,
+        # ---- Loss: InfoNCELoss-only ----
+        tau_final: float = 0.01,                 # cosine-schedule final tau (when not learnable)
+        weight_q2g: float = 0.6,
+        weight_g2q: float = 0.4,
+        hard_mining_k: int = 0,
+        neg_bank_size: int = 0,
+        # ---- Loss: WeightedInfoNCELoss-only ----
+        weighted_loss_k: float = 5.0,            # sigmoid steepness for weight→eps mapping
+        # ---- Optimizer ----
+        learning_rate: float = 1e-4,
+        text_lr_factor: float = 0.1,             # lr * factor for DGTRS-CLIP/LoRA params
+        weight_decay: float = 1e-4,
+        grad_clip: float = 1.0,
+        # ---- Sampler ----
+        sampler_type: str = "mutex",             # 'mutex' | 'dss' | 'none'
+        dss_warmup_epochs: int = 1,
+        dss_reembed_every: int = 1,
+        dss_knn_device: str = "cuda",
+        dss_use_lsh: bool = False,
+        dss_lsh_num_tables: int = 8,
+        dss_lsh_num_bits: int = 14,
+        dss_cache_dir: str | None = None,
+        # Legacy alias (kept until train_gtauav.py is rewritten in step 4).
+        use_mutex_sampler: bool = True,
+    ) -> None:
+        # Loss (shared).
+        self.loss_type = loss_type
+        self.tau_init = tau_init
+        self.tau_min = tau_min
+        self.tau_max = tau_max
+        self.learnable_temperature = learnable_temperature
+        self.label_smoothing = label_smoothing
+        # Loss (InfoNCE-specific).
+        self.tau_final = tau_final
+        self.weight_q2g = weight_q2g
+        self.weight_g2q = weight_g2q
+        self.hard_mining_k = hard_mining_k
+        self.neg_bank_size = neg_bank_size
+        # Loss (WeightedInfoNCE-specific).
+        self.weighted_loss_k = weighted_loss_k
+        # Optimizer.
+        self.learning_rate = learning_rate
+        self.text_lr_factor = text_lr_factor
+        self.weight_decay = weight_decay
+        self.grad_clip = grad_clip
+        # Sampler.
+        self.sampler_type = sampler_type
+        self.dss_warmup_epochs = dss_warmup_epochs
+        self.dss_reembed_every = dss_reembed_every
+        self.dss_knn_device = dss_knn_device
+        self.dss_use_lsh = dss_use_lsh
+        self.dss_lsh_num_tables = dss_lsh_num_tables
+        self.dss_lsh_num_bits = dss_lsh_num_bits
+        self.dss_cache_dir = dss_cache_dir
+        self.use_mutex_sampler = use_mutex_sampler
+
+
+def get_training_cfg(path2cfg: str) -> TrainingConfig:
+    """Load ONLY training config (TESTING ONLY — production uses load_all_configs)."""
+    gin.clear_config()
+    gin.parse_config_file(f"{path2cfg}training.gin")
+    return TrainingConfig()
+
+

src/datasets/gtauav_dataset.py

@@ -30,13 +30,17 @@ LOGGER = logging.getLogger("caption_test.gtauav_dataset")
 coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
 
 # Default paths.
-_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
-_CAPTION_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-captions")
+_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
+_CAPTION_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions")
 _EMPTY_CAPTION = ""
 
 # Regex to split P1/P2/P3 sections.
 _P_SPLIT = re.compile(r"\*\*P[123][^*]*\*\*\s*:?\s*")
 
+# TODO: Transforms
+# sat_transforms = ...
+# drone_transforms = ...
+
 
 def _parse_caption_levels(output: str) -> tuple[str, str, str]:
     """Split VLM caption output into L1, L2, L3 levels.

src/eval/evaluator.py

@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+"""Retrieval evaluation for GTA-UAV-LR cross-view geo-localization.
+
+Computes R@K and MRR for both q→g (drone→satellite) and g→q (satellite→drone)
+on the full satellite gallery. Multi-match: a query counts as a hit@K if ANY
+of its valid satellite matches (sat_candidates) appears in the top-K.
+
+Body transplanted byte-for-byte from src/training/train_gtauav.py::_evaluate
+in the main branch. The single difference is the type annotation
+`model: AsymmetricEncoder` → `model: nn.Module` (relaxed for duck-typing
+across encoder families); semantically identical to the main-branch version.
+
+Note: not to be confused with src/eval/evaluate.py (legacy v2 helper for
+UAV-VisLoc with a different signature). This module lives at
+src/eval/evaluator.py and is the active evaluator for v3 GTA-UAV-LR.
+"""
+
+import logging
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from src.models.asymmetric_encoder import AsymmetricEncoder
+from src.datasets.gtauav_dataset import (
+    GTAUAVDataset,
+    GTAUAVDroneQuery,
+    GTAUAVSatGallery,
+    collate_drone_query,
+    collate_sat_gallery,
+)
+
+LOGGER = logging.getLogger("caption_test.evaluator")
+
+
+@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.
+
+    Args:
+        model: Encoder with `encode_query(drone_img, l1, l2, l3)`
+               and `encode_gallery(sat_img, l1, l2, l3)`. Must expose
+               `fusion_query.gate_value` and `fusion_gallery.gate_value`.
+        loader: DataLoader over a GTAUAVDataset (used only to pull dataset
+                + batch_size/num_workers/pin_memory; iteration is bypassed —
+                we build separate query and gallery loaders inside).
+        device: Torch device string.
+        loss_fn: If provided, computes per-batch loss against paired gallery
+                 entries (uses the first valid sat per query as its positive).
+                 The mean loss appears in the returned dict under 'loss'.
+        epoch, total_epochs: Passed through to loss_fn.
+        k_values: K values for R@K (e.g. (1, 5, 10)).
+        max_batches: Cap on query batches for quick sanity checks (gallery
+                     is always full).
+        desc: tqdm description prefix.
+
+    Returns:
+        Dict with: r@K_q2g, ap_q2g (= MRR), r@K_g2q, ap_g2q, loss (optional),
+        n_query, n_gallery, n_scored_g2q, gate_q, gate_g.
+    """
+    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)
+        q = model.encode_query(
+            drone_img,
+            batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
+        )
+        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
+
+

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).

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).

src/main.py

@@ -0,0 +1,73 @@
+"""Entry point: load configs and run training.
+
+Usage: (commands for each preset)
+    python -m src.main gtauav_balanced
+    python -m src.main gtauav_balanced_asym
+    python -m src.main gtauav_balanced_stripnet
+    python -m src.main gtauav_balanced_stripnet_unfrozen
+    python -m src.main gtauav_baseline
+    python -m src.main gtauav_baseline_asym
+    python -m src.main gtauav_baseline_stripnet
+    python -m src.main gtauav_baseline_stripnet_unfrozen
+    python -m src.main gtauav_image_heavy
+    python -m src.main gtauav_gtauav_text_heavy
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+
+import coloredlogs
+
+from src.conf.config_loader import load_all_configs
+from src.training.trainer_new import Trainer
+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 <preset_name>\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
+
+    # -------------------------------------------------------
+    ''' ONLY FOR DEBUG with launch.json config:
+         "args": ["main gtauav_balanced"] -> so need to extract
+         preset name "gtauav_balanced"
+    '''
+    preset_name = preset_name.split(' ')[1]
+    # -------------------------------------------------------
+
+    configs = load_all_configs(path2cfg, preset_name)
+
+    trainer = Trainer(
+        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"],
+    )
+
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()
+
+

src/models/asymmetric_encoder.py

@@ -192,23 +192,25 @@ class DINOv3ViT(nn.Module):
 
     @classmethod
     def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
-        """Load from .pth or .safetensors checkpoint."""
-        model = cls()
-        path = Path(path)
-        LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
-        if path.suffix == ".safetensors":
-            state = load_safetensors(str(path))
-        else:
-            state = torch.load(str(path), map_location="cpu", weights_only=False)
-            if "model" in state:
-                state = state["model"]
-            elif "state_dict" in state:
-                state = state["state_dict"]
-        model.load_state_dict(state, strict=False)
-        n_params = sum(p.numel() for p in model.parameters())
-        LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
-        return model
-
+        try:
+            """Load from .pth or .safetensors checkpoint."""
+            model = cls()
+            path = Path(path)
+            LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
+            if path.suffix == ".safetensors":
+                state = load_safetensors(str(path))
+            else:
+                state = torch.load(str(path), map_location="cpu", weights_only=False)
+                if "model" in state:
+                    state = state["model"]
+                elif "state_dict" in state:
+                    state = state["state_dict"]
+            model.load_state_dict(state, strict=False)
+            n_params = sum(p.numel() for p in model.parameters())
+            LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
+            return model
+        except FileNotFoundError as e:
+            LOGGER.exception(msg=e.strerror)
 
 # LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
 # in src/models/dgtrs/model.py (DGTRSTextEncoder)

src/training/csv_logger.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+"""Per-batch and per-epoch CSV logger.
+
+Writes:
+    {output_dir}/logs/train.csv               — epoch-level train averages
+    {output_dir}/logs/val.csv                  — epoch-level val metrics
+    {output_dir}/logs/train_recall.csv         — epoch-level train recall metrics
+    {output_dir}/logs/train_batches.csv        — per-batch train metrics (all epochs)
+    {output_dir}/logs/epoch_{N}_batches.csv    — per-batch for one epoch
+
+Body transplanted verbatim from src/training/train_gtauav.py (pre-step-4b)
+with no logic changes — only the relocation.
+"""
+
+import logging
+from pathlib import Path
+
+import pandas as pd
+
+LOGGER = logging.getLogger("caption_test.csv_logger")
+
+
+class CSVLogger:
+    """Log train/val metrics to CSV files using pandas."""
+
+    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,
+        )
+
+        

src/training/train.py

@@ -1,273 +0,0 @@
-from __future__ import annotations
-
-"""Training loop for caption quality test on cross-view geo-localization.
-
-GeoRSCLIP dual encoder with GatedFusion on query branch.
-Single InfoNCE loss: query(drone+text) vs gallery(satellite).
-"""
-
-import argparse
-import json
-import logging
-import time
-from pathlib import Path
-
-import gin
-import torch
-import torch.nn as nn
-from torch.amp import GradScaler, autocast
-from torch.optim import AdamW
-from torch.optim.lr_scheduler import CosineAnnealingLR
-from torch.utils.data import DataLoader
-
-from src.datasets.visloc_with_captions import (
-    GeoLocCaptionDataset,
-    collate_caption_batch,
-)
-from src.eval.evaluate import evaluate_retrieval
-from src.losses.multi_infonce import InfoNCELoss
-from src.models.dual_encoder import DualEncoderCaptionTest
-
-LOGGER = logging.getLogger("caption_test.train")
-
-
-@gin.configurable
-class TrainConfig:
-    """Top-level training configuration.
-
-    Args:
-        train_query_file: Path to train_query.txt.
-        val_query_file: Path to test_query.txt (used as val).
-        data_root: Root of UAV-GeoLoc dataset.
-        output_dir: Checkpoint and log output directory.
-        epochs: Number of training epochs.
-        batch_size: Mini-batch size.
-        num_workers: DataLoader workers.
-        learning_rate: AdamW initial LR.
-        weight_decay: AdamW weight decay.
-        grad_clip: Max gradient norm (0 disables).
-        use_amp: Enable fp16 mixed-precision.
-        eval_every: Run validation every N epochs.
-        seed: Random seed.
-        device: torch device.
-    """
-
-    def __init__(
-        self,
-        train_query_file: str = "Index/train_query.txt",
-        val_query_file: str = "Index/test_query.txt",
-        data_root: str = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc",
-        output_dir: str = "out/caption_test",
-        epochs: int = 10,
-        batch_size: int = 128,
-        num_workers: int = 4,
-        learning_rate: float = 1e-4,
-        weight_decay: float = 1e-4,
-        grad_clip: float = 1.0,
-        use_amp: bool = True,
-        eval_every: int = 2,
-        seed: int = 42,
-        device: str = "cuda",
-    ) -> None:
-        self.train_query_file = train_query_file
-        self.val_query_file = val_query_file
-        self.data_root = data_root
-        self.output_dir = Path(output_dir)
-        self.epochs = epochs
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.learning_rate = learning_rate
-        self.weight_decay = weight_decay
-        self.grad_clip = grad_clip
-        self.use_amp = use_amp
-        self.eval_every = eval_every
-        self.seed = seed
-        self.device = device
-
-
-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 train(config_path: str) -> None:
-    """Run full training loop from gin config."""
-    gin.parse_config_file(config_path)
-    cfg = TrainConfig()
-
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s %(name)s %(levelname)s %(message)s",
-    )
-    _set_seed(cfg.seed)
-    cfg.output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Model + loss.
-    model = DualEncoderCaptionTest().to(cfg.device)
-    loss_fn = InfoNCELoss().to(cfg.device)
-
-    preprocess = model.preprocess
-
-    train_ds = GeoLocCaptionDataset(
-        query_file=cfg.train_query_file,
-        data_root=cfg.data_root,
-        image_transform=preprocess,
-    )
-    val_ds = GeoLocCaptionDataset(
-        query_file=cfg.val_query_file,
-        data_root=cfg.data_root,
-        image_transform=preprocess,
-    )
-
-    train_loader = DataLoader(
-        train_ds,
-        batch_size=cfg.batch_size,
-        shuffle=True,
-        num_workers=cfg.num_workers,
-        collate_fn=collate_caption_batch,
-        pin_memory=True,
-        drop_last=True,
-    )
-    val_loader = DataLoader(
-        val_ds,
-        batch_size=cfg.batch_size,
-        shuffle=False,
-        num_workers=cfg.num_workers,
-        collate_fn=collate_caption_batch,
-        pin_memory=True,
-    )
-
-    optimizer = AdamW(
-        model.trainable_parameters(),
-        lr=cfg.learning_rate,
-        weight_decay=cfg.weight_decay,
-    )
-    scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs)
-    scaler = GradScaler(enabled=cfg.use_amp)
-
-    n_trainable = sum(p.numel() for p in model.trainable_parameters())
-    n_total = sum(p.numel() for p in model.parameters())
-    LOGGER.info(
-        "trainable=%d (%.2f%%) total=%d train=%d val=%d",
-        n_trainable, 100.0 * n_trainable / n_total,
-        n_total, len(train_ds), len(val_ds),
-    )
-
-    history: list[dict] = []
-
-    for epoch in range(cfg.epochs):
-        model.train()
-        epoch_start = time.time()
-        agg: dict[str, float] = {}
-        n_batches = 0
-
-        for batch in train_loader:
-            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)
-            caption_drone = batch["caption_drone"]
-
-            with autocast(device_type="cuda", enabled=cfg.use_amp):
-                embeddings = model(
-                    drone_img=drone_img,
-                    sat_img=sat_img,
-                    caption_drone=caption_drone,
-                )
-                loss_dict = loss_fn(
-                    embeddings=embeddings,
-                    epoch=epoch,
-                    total_epochs=cfg.epochs,
-                )
-
-            total_loss = loss_dict["total"]
-            scaler.scale(total_loss).backward()
-
-            if cfg.grad_clip > 0:
-                scaler.unscale_(optimizer)
-                nn.utils.clip_grad_norm_(
-                    model.trainable_parameters(),
-                    max_norm=cfg.grad_clip,
-                )
-            scaler.step(optimizer)
-            scaler.update()
-
-            for key, val in loss_dict.items():
-                agg[key] = agg.get(key, 0.0) + float(val.item())
-            n_batches += 1
-
-        scheduler.step()
-        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=%.4f",
-            epoch, elapsed,
-            optimizer.param_groups[0]["lr"],
-            means.get("total", 0.0),
-            means.get("temperature", 0.0),
-            means.get("gate", 1.0),
-        )
-
-        epoch_record: dict = {
-            "epoch": epoch,
-            "elapsed_seconds": elapsed,
-            "train": means,
-        }
-
-        # Validation.
-        if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
-            model.eval()
-            val_metrics = evaluate_retrieval(
-                model=model,
-                loader=val_loader,
-                device=cfg.device,
-            )
-            epoch_record["val"] = val_metrics
-            LOGGER.info(
-                "val epoch=%d R@1_q2g=%.4f R@5_q2g=%.4f R@10_q2g=%.4f",
-                epoch,
-                val_metrics.get("r@1_query_to_gallery", 0.0),
-                val_metrics.get("r@5_query_to_gallery", 0.0),
-                val_metrics.get("r@10_query_to_gallery", 0.0),
-            )
-
-        history.append(epoch_record)
-
-        _atomic_save(
-            obj={
-                "epoch": epoch,
-                "model_state": model.state_dict(),
-                "optimizer_state": optimizer.state_dict(),
-                "config_path": config_path,
-            },
-            path=cfg.output_dir / f"ckpt_epoch{epoch:03d}.pt",
-        )
-
-    history_path = cfg.output_dir / "history.json"
-    with history_path.open("w", encoding="utf-8") as f:
-        json.dump(history, f, indent=2)
-
-    LOGGER.info("training complete, history saved to %s", history_path)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Caption quality test training.")
-    parser.add_argument("--config", type=str, required=True, help="Gin config file.")
-    args = parser.parse_args()
-    train(config_path=args.config)
-
-
-if __name__ == "__main__":
-    main()

src/utils/__init__.py

@@ -0,0 +1,2 @@
+"""Utilities: paths, seeding, IO."""
+

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)
+
+
+
+        

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}",
+    )
+

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)
+
+
+    

tests/conftest.py

@@ -0,0 +1,157 @@
+"""Shared pytest fixtures + reporter hooks for caption-test test suite.
+
+Provides:
+  - proj_dir, path2cfg: real repository paths for integration-style tests
+    that load actual gin presets from in/config_files/.
+  - clear_gin: autouse fixture that wipes gin global state before every test
+    (gin keeps bindings in module-level singleton; tests must not leak).
+
+Reporter hooks print "✅/❌ <docstring>" next to each test's PASSED/FAILED
+line in -v mode.
+
+Preset name lists live in a separate module (`tests/_presets.py`) so test
+files can import them with a plain `import _presets` — no relative imports,
+no need for tests/ to be a package.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import gin
+import pytest
+
+
+DINOV3_PRESETS = (
+    "gtauav_balanced",
+    "gtauav_balanced_asym",
+    "gtauav_baseline",
+    "gtauav_baseline_asym",
+    "gtauav_image_heavy",
+    "gtauav_text_heavy",
+)
+ 
+STRIPNET_PRESETS = (
+    "gtauav_balanced_stripnet",
+    "gtauav_balanced_stripnet_unfrozen",
+    "gtauav_baseline_stripnet",
+    "gtauav_baseline_stripnet_unfrozen",
+)
+ 
+ALL_TRAINING_PRESETS = DINOV3_PRESETS + STRIPNET_PRESETS
+
+# --- gin hygiene -----------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def clear_gin():
+    """Wipe gin's global binding state before AND after each test.
+
+    Gin keeps bindings in module-level singletons; without this fixture, a
+    test that loads a config (or even just calls gin.parse_config_file in a
+    helper) leaks bindings into the next test, leading to flaky failures
+    like 'Unknown configurable' or wrong field values.
+    """
+    gin.clear_config()
+    yield
+    gin.clear_config()
+
+
+# --- real-repo paths -------------------------------------------------------
+
+
+@pytest.fixture
+def proj_dir() -> Path:
+    """Path to repository root (the directory that contains src/, tests/, in/)."""
+    return Path(__file__).resolve().parent.parent
+
+
+@pytest.fixture
+def path2cfg(proj_dir: Path) -> str:
+    """Trailing-slashed path to in/config_files/, matching `src/main.py`.
+
+    Per REQUIREMENTS_GIN_STYLE.md §5, src/main.py builds this path as
+    `f"{proj_dir}in/config_files/"`. Tests that exercise the real repo
+    layout should use this fixture verbatim instead of constructing it
+    independently.
+    """
+    return f"{proj_dir}/in/config_files/"
+
+
+# --- reporter hooks --------------------------------------------------------
+
+
+def _docstring_summary(item: pytest.Item) -> str | None:
+    """Return the first non-empty line of a test's docstring, or None."""
+    func = getattr(item, "function", None) or getattr(item, "obj", None)
+    if func is None or not getattr(func, "__doc__", None):
+        return None
+    for line in func.__doc__.strip().splitlines():
+        stripped = line.strip()
+        if stripped:
+            return stripped
+    return None
+
+
+# Cache nodeid → docstring summary, populated at collection time so the
+# logreport hook can look them up without re-introspecting the test function.
+_LAST_SEEN_SUMMARY: dict[str, str] = {}
+
+
+def pytest_collection_modifyitems(
+    config: pytest.Config,
+    items: list[pytest.Item],
+) -> None:
+    """Cache each item's docstring summary for later use by the status hook."""
+    for item in items:
+        summary = _docstring_summary(item)
+        if summary:
+            _LAST_SEEN_SUMMARY[item.nodeid] = summary
+
+
+def pytest_runtest_logreport(report: pytest.TestReport) -> None:
+    """Print test results with parametrized tests grouped under one header.
+
+    Non-parametrized test:
+        ✅  <docstring summary>
+
+    Parametrized test (first occurrence of the group):
+        <docstring summary>
+          ✅  <param>
+
+    Parametrized test (subsequent occurrences):
+          ✅  <param>
+
+    Pytest emits 3 reports per test (setup → call → teardown). We hook the
+    `call` phase — the one where pass/fail is actually decided.
+    """
+    if report.when != "call":
+        return
+
+    summary = _LAST_SEEN_SUMMARY.get(report.nodeid, "(no docstring)")
+    if report.passed:
+        icon = "✅"
+    elif report.failed:
+        icon = "❌"
+    else:
+        icon = "⏭️"
+
+    # Detect parametrization: pytest encodes params as `nodeid[param1-param2-...]`.
+    if "[" in report.nodeid and report.nodeid.endswith("]"):
+        base_id, _, param_part = report.nodeid.partition("[")
+        param_label = param_part[:-1] or "empty" # strip trailing ']'
+
+        # Print docstring header only on first encounter of this group.
+        # Leading "\n" separates the header from pytest's progress dot.
+        if base_id not in _PRINTED_GROUP_HEADERS:
+            print(f"\n{summary}")
+            _PRINTED_GROUP_HEADERS.add(base_id)
+        print(f"  {icon}  {param_label}")
+    else:
+        print(f"  {icon}  {summary}")
+
+
+# Tracks which parametrized test groups have already had their docstring
+# header printed. Reset implicitly at the start of each pytest run because
+# Python re-imports conftest.py.
+_PRINTED_GROUP_HEADERS: set[str] = set()

tests/test_trainer.py

@@ -0,0 +1,187 @@
+"""Tests for src.training.trainer_new.Trainer.
+
+Scope: __init__ behaviour, backbone validation, ModelsConfig type union.
+Out of scope: actual training (requires GPU + datasets + model checkpoints).
+
+The Trainer class is designed to defer all heavy lifting (CUDA, model
+construction, dataset loading) to .train(); __init__ just stores the 6 cfg
+objects and zeros out runtime state. This makes it cheap to test.
+"""
+
+from __future__ import annotations
+
+from typing import get_args
+
+import pytest
+
+from src.conf.config_loader import load_all_configs
+from src.conf.models_dinov3_conf import DINOv3ModelsConfig
+from src.conf.models_stripnet_conf import StripNetModelsConfig
+from src.training.trainer_new import (
+    ModelsConfig,
+    Trainer,
+    _SUPPORTED_BACKBONES,
+)
+
+from conftest import DINOV3_PRESETS, STRIPNET_PRESETS
+
+
+# -- module-level constants ------------------------------------------------
+
+
+def test_supported_backbones_is_frozenset() -> None:
+    """_SUPPORTED_BACKBONES must be a frozenset (immutable, hashable)."""
+    assert isinstance(_SUPPORTED_BACKBONES, frozenset)
+
+
+def test_supported_backbones_contents() -> None:
+    """Exactly dinov3 and stripnet are supported in the current refactor.
+
+    Sofia (v1/v71) is intentionally absent — see Trainer._validate_backbone
+    for the rationale and the steps to add it later.
+    """
+    assert _SUPPORTED_BACKBONES == frozenset({"dinov3", "stripnet"})
+
+
+def test_models_config_union_contents() -> None:
+    """ModelsConfig union mirrors _SUPPORTED_BACKBONES (dinov3 | stripnet)."""
+    union_members = set(get_args(ModelsConfig))
+    assert union_members == {DINOv3ModelsConfig, StripNetModelsConfig}
+
+
+# -- _validate_backbone --------------------------------------------------
+
+
+@pytest.mark.parametrize("backbone", ["dinov3", "stripnet"])
+def test_validate_backbone_accepts_supported(
+    path2cfg: str, backbone: str,
+) -> None:
+    """Supported backbones pass _validate_backbone silently.
+
+    We use a real preset to build a valid Trainer — this also exercises
+    the load_all_configs → Trainer(...) integration.
+    """
+    preset = "gtauav_balanced" if backbone == "dinov3" else "gtauav_balanced_stripnet"
+    cfgs = load_all_configs(path2cfg, preset)
+    trainer = Trainer(
+        pipeline_cfg=cfgs["pipeline"],
+        hardware_cfg=cfgs["hardware"],
+        training_cfg=cfgs["training"],
+        tracking_cfg=cfgs["tracking"],
+        models_common_cfg=cfgs["models_common"],
+        models_cfg=cfgs["models"],
+    )
+    # Must not raise.
+    trainer._validate_backbone()
+
+
+@pytest.mark.parametrize("bad_backbone", ["sofia_v1", "sofia_v71", "mistral_42b", ""])
+def test_validate_backbone_rejects_unsupported(
+    path2cfg: str, bad_backbone: str,
+) -> None:
+    """Unsupported backbones (incl. sofia) raise NotImplementedError, not ImportError.
+
+    The user must get a clear, actionable message — not a stack trace from
+    a missing module.
+    """
+    cfgs = load_all_configs(path2cfg, "gtauav_balanced")
+    trainer = Trainer(
+        pipeline_cfg=cfgs["pipeline"],
+        hardware_cfg=cfgs["hardware"],
+        training_cfg=cfgs["training"],
+        tracking_cfg=cfgs["tracking"],
+        models_common_cfg=cfgs["models_common"],
+        models_cfg=cfgs["models"],
+    )
+    # Tamper with backbone — simulate what would happen if config_loader
+    # were extended to accept sofia presets.
+    trainer.models_common_cfg.backbone = bad_backbone
+
+    with pytest.raises(NotImplementedError) as excinfo:
+        trainer._validate_backbone()
+
+    # Error message must mention the offending backbone name and what's supported.
+    msg = str(excinfo.value)
+    assert bad_backbone in msg or repr(bad_backbone) in msg
+    assert "dinov3" in msg
+    assert "stripnet" in msg
+
+
+# -- Trainer.__init__ smoke tests ------------------------------------------
+
+
+@pytest.mark.parametrize("preset_name", DINOV3_PRESETS + STRIPNET_PRESETS)
+def test_trainer_init_with_real_preset(path2cfg: str, preset_name: str) -> None:
+    """Trainer(...) instantiates from every real preset's loaded cfgs.
+
+    Heavy work (CUDA, model build, dataset open) is deferred to .train();
+    __init__ only stores cfgs and zeros runtime state, so this is cheap and
+    GPU-free.
+    """
+    cfgs = load_all_configs(path2cfg, preset_name)
+
+    trainer = Trainer(
+        pipeline_cfg=cfgs["pipeline"],
+        hardware_cfg=cfgs["hardware"],
+        training_cfg=cfgs["training"],
+        tracking_cfg=cfgs["tracking"],
+        models_common_cfg=cfgs["models_common"],
+        models_cfg=cfgs["models"],
+    )
+
+    # Cfgs are stored as-is.
+    assert trainer.pipeline_cfg is cfgs["pipeline"]
+    assert trainer.hardware_cfg is cfgs["hardware"]
+    assert trainer.training_cfg is cfgs["training"]
+    assert trainer.tracking_cfg is cfgs["tracking"]
+    assert trainer.models_common_cfg is cfgs["models_common"]
+    assert trainer.models_cfg is cfgs["models"]
+
+
+def test_trainer_init_zeros_runtime_state(path2cfg: str) -> None:
+    """All runtime fields are None / 0 / [] before .train() is called."""
+    cfgs = load_all_configs(path2cfg, "gtauav_balanced")
+    trainer = Trainer(
+        pipeline_cfg=cfgs["pipeline"],
+        hardware_cfg=cfgs["hardware"],
+        training_cfg=cfgs["training"],
+        tracking_cfg=cfgs["tracking"],
+        models_common_cfg=cfgs["models_common"],
+        models_cfg=cfgs["models"],
+    )
+
+    # None-typed runtime fields.
+    for attr in (
+        "output_dir", "full_config", "tracker", "csv_logger", "model",
+        "loss_fn", "neg_bank", "optimizer", "scheduler", "scaler",
+        "train_ds", "test_ds", "train_eval_ds",
+        "train_loader", "test_loader", "train_eval_loader",
+        "batch_sampler", "emb_cache", "profiler", "resume_ckpt",
+    ):
+        assert getattr(trainer, attr) is None, (
+            f"trainer.{attr} should be None before .train(), "
+            f"got {type(getattr(trainer, attr)).__name__}"
+        )
+
+    # Counter / loop state initialized to identity values.
+    assert trainer.start_epoch == 0
+    assert trainer.global_step == 0
+    assert trainer.best_r1 == 0.0
+    assert trainer.history == []
+    assert trainer.steps_per_epoch == 0
+
+
+# -- Trainer.train end-to-end signature ------------------------------------
+
+
+def test_trainer_train_method_exists_and_takes_no_args() -> None:
+    """Trainer.train() takes only `self` — main.py calls trainer.train()."""
+    import inspect
+
+    sig = inspect.signature(Trainer.train)
+    params = [p for p in sig.parameters.values() if p.name != "self"]
+    assert params == [], (
+        f"Trainer.train() must take only self; got extra params: {params}"
+    )
+
+