caption-test/refactor_proposal_v2.md

# Рефакторинг `caption-test` (`belka_refactor`) — обновлённое предложение

> **Версия 2.** Обновлена под полный набор требований: `REQUIREMENTS_GIN_STYLE.md` (включая центральный `load_all_configs`), стандарт `code-style`, рекомендации «5 конфигов — оптимум» с принципом разделения «меняются вместе → один конфиг».

---

## 0. Что изменилось относительно версии 1

| # | Было в v1 | Стало в v2 | Почему |
|---|---|---|---|
| 1 | 6 конфиг-классов | **5 конфиг-классов** (Pipeline, Hardware, Models, Training, Tracking) | «5 — оптимальное число»; принцип «оси изменчивости»: Loss+Optimizer+Sampling меняются вместе при экспериментах с обучением → один `TrainingConfig` |
| 2 | 6 индивидуальных `get_*_cfg()` в `main()` | **Один** `load_all_configs()` + `gin.clear_config()` внутри | Прямое требование `REQUIREMENTS_GIN_STYLE.md §3.1`: индивидуальные loader'ы — **только для тестов** |
| 3 | `os.environ.get("CAPTION_TEST_PRESET")` для выбора пресета | `path2cfg` фиксирован, разные пресеты = разные директории | env vars — это скрытый CLI; стандарт требует «всё в `.gin`» |
| 4 | Подкаталог `presets/` внутри `in/config_files/` | `in/config_files/` напрямую (один пресет = один запуск; копировать дир для другого) | Соответствие референсной структуре `test_bb_uav` |
| 5 | Не подсвечен антипаттерн `@gin.configurable` на `InfoNCELoss` + дубль биндингов | **Явно подсвечен**: один источник правды | Двойная gin-регистрация → тихие баги |
| 6 | `get_proj_dir()` упомянут абстрактно | Конкретная реализация через MARKERS (`pyproject.toml`, `.git`, `in`) | Из `REQUIREMENTS_GIN_STYLE.md §5` |

---

## 1. Главные нарушения стандарта в текущем коде

### 1.1 Критические (запрещены прямо)

```python
# src/training/train_gtauav.py, строки ~80
@gin.configurable(module="src.training.train_gtauav")
@dataclass                                              # ← FORBIDDEN
class TrainConfigGTAUAV:
    train_json: str = _TRAIN_JSON
    # ... 50+ полей
```

> «Запрещено использовать `dataclass` совместно с gin» — стандарт §3.1 и `Reference Examples → Anti-patterns`.

```python
# src/training/train_gtauav.py, ~main()
def main() -> None:
    parser = argparse.ArgumentParser(...)              # ← FORBIDDEN
    parser.add_argument("--config", ...)
    parser.add_argument("--baseline", ...)
    # ... 15+ CLI флагов
```

> «Запрещён argparse — все параметры из .gin файлов» — стандарт §3.4.

```python
# src/training/train_gtauav.py, module level
_RGB_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR"        # hardcoded
_DINO_WEB = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth"
# ... ещё 5 hardcoded путей
```

> «Нет захардкоженных model ID / промптов / размеров» — чеклист §6.

### 1.2 Серьёзные структурные

- **Один мега-конфиг на 50+ полей** — нарушает «1 .gin = 1 конфиг-класс».
- **Дубль биндингов в `gtauav_balanced.gin`**: одни и те же значения прописаны и для `TrainConfigGTAUAV.tau_init`, и для `InfoNCELoss.temperature_init`. Если кто-то поменяет одно и забудет другое — тихий рассинхрон.
- **`InfoNCELoss` с `@gin.configurable`** + одновременно создаётся через явные kwargs из `cfg.tau_init`. Двойной источник правды.
- **`@torch.no_grad()` вместо `@torch.inference_mode()`** в `_evaluate`, `_embed_drone_queries`.
- **`train_gtauav.py` — 1296 строк**: dataclass-конфиг + утилиты + CSVLogger + `_evaluate` (150 строк) + `train()` (700 строк) + `main()` с argparse.

---

## 2. Целевая архитектура — 5 конфигов

### Принцип разделения «оси изменчивости»

| Что меняется вместе | Конфиг |
|---|---|
| Запуск пайплайна (что обрабатываем, куда сохраняем, сколько эпох, расписание eval) | **PipelineConfig** |
| GPU / память / производительность (`batch_size`, `num_workers`, `use_amp`, `grad_checkpointing`) | **HardwareConfig** |
| Архитектура (бэкбоны, пути к чекпоинтам, MONA, gate, baseline_mode, asym/shared) | **ModelsConfig** |
| Обучение (loss + optimizer + sampler — экспериментируем с этим вместе) | **TrainingConfig** |
| Наблюдаемость (W&B, TB, Grad-CAM, profiler) — независимый рубильник | **TrackingConfig** |

> **Почему Loss + Optimizer + Sampling в одном `TrainingConfig`?** Они меняются вместе. Когда исследуешь «что улучшит R@1» — ты крутишь `tau`, `label_smoothing`, `lr`, `text_lr_factor`, `sampler_type`, `dss_warmup`. Они образуют единый «recipe обучения». Разделение их на 3 конфига создавало бы координационную нагрузку без выгоды.
>
> **Почему Tracking отдельно?** Это диагностика. Можно запустить один и тот же эксперимент с `use_wandb=True` и без — результат обучения идентичный, меняется только что мы записываем. Это **независимая ось**, поэтому отдельный конфиг.

### Целевая структура каталогов

```
caption-test/
├── in/
│   └── config_files/                       # Один активный пресет
│       ├── pipeline.gin                    # 1 файл = 1 конфиг-класс
│       ├── hardware.gin
│       ├── models.gin
│       ├── training.gin
│       └── tracking.gin
├── presets/                                # Готовые пресеты (копируются в in/config_files/)
│   ├── gtauav_balanced/
│   │   ├── pipeline.gin
│   │   ├── hardware.gin
│   │   ├── models.gin
│   │   ├── training.gin
│   │   └── tracking.gin
│   ├── gtauav_baseline/
│   ├── gtauav_balanced_asym/
│   ├── gtauav_balanced_stripnet/
│   ├── gtauav_balanced_stripnet_unfrozen/
│   ├── gtauav_text_heavy/
│   └── gtauav_image_heavy/
├── src/
│   ├── conf/
│   │   ├── __init__.py
│   │   ├── pipeline_conf.py                # PipelineConfig + get_pipeline_cfg
│   │   ├── hardware_conf.py                # HardwareConfig + get_hardware_cfg
│   │   ├── models_conf.py                  # ModelsConfig + get_models_cfg
│   │   ├── training_conf.py                # TrainingConfig + get_training_cfg
│   │   ├── tracking_conf.py                # TrackingConfig + get_tracking_cfg
│   │   └── config_loader.py                # load_all_configs() — единый продакшен-вход
│   ├── datasets/                           # БЕЗ ИЗМЕНЕНИЙ (логика самплеров корректна)
│   ├── models/                             # БЕЗ ИЗМЕНЕНИЙ (модельная архитектура работает)
│   ├── losses/                             # минимальные правки (см. §4)
│   ├── eval/
│   │   ├── __init__.py
│   │   └── evaluator.py                    # NEW — _evaluate() сюда, @torch.inference_mode()
│   ├── training/
│   │   ├── __init__.py
│   │   ├── trainer.py                      # NEW — class Trainer: основной цикл
│   │   ├── csv_logger.py                   # NEW — CSVLogger вытащен сюда
│   │   ├── trackers.py                     # без изменений
│   │   ├── grad_monitor.py                 # без изменений
│   │   ├── gradcam.py                      # без изменений
│   │   ├── profiling.py                    # без изменений
│   │   └── plot_metrics.py                 # без изменений
│   ├── utils/
│   │   ├── __init__.py
│   │   ├── io_utils.py                     # atomic_save_torch, clear_vram
│   │   ├── seed_utils.py                   # set_seed
│   │   └── path_utils.py                   # get_proj_dir
│   └── main.py                             # NEW — единственная точка входа
├── scripts/                                # без изменений (это уже отдельные tools)
├── meta/                                   # без изменений (data artifacts)
└── nn_models/                              # без изменений (checkpoints, gitignored)
```

**Workflow смены пресета:**

```bash
# Вместо CLI флагов — копируем нужный пресет:
cp -r presets/gtauav_baseline/* in/config_files/
python -m src.main

# Или для эксперимента: создать новый пресет
cp -r presets/gtauav_balanced presets/gtauav_balanced_lr5e5
# отредактировать presets/gtauav_balanced_lr5e5/training.gin
cp -r presets/gtauav_balanced_lr5e5/* in/config_files/
python -m src.main
```

Каждый запуск однозначно описан содержимым `in/config_files/`. Воспроизводимость = снэпшот этой директории.

---

## 3. Конфиг-классы — полный код

### 3.1 `src/conf/pipeline_conf.py`

```python
from __future__ import annotations

import gin


@gin.configurable
class PipelineConfig:
    """Pipeline orchestration: data IO, training schedule, output, resume."""

    def __init__(
        self,
        # Data paths.
        train_json: str = "meta/train_80.json",
        test_json: str = "meta/test_20.json",
        rgb_root: str = "/data/GTA-UAV-LR",
        caption_root: str = "/data/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 — use load_all_configs in production).

    Args:
        path2cfg: Path to config directory (with trailing slash).

    Returns:
        Instantiated PipelineConfig.
    """
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}pipeline.gin")
    return PipelineConfig()
```

### 3.2 `src/conf/hardware_conf.py`

```python
from __future__ import annotations

import gin


@gin.configurable
class HardwareConfig:
    """GPU profile + memory/compute optimisation flags.

    Everything that changes when you switch hardware (4090 → A100 → Jetson)
    lives here. batch_size and grad_accum_steps are hardware-bound: they
    determine VRAM footprint, not the training recipe.
    """

    def __init__(
        self,
        device: str = "cuda",
        batch_size: int = 8,
        grad_accum_steps: int = 1,
        num_workers: int = 4,
        use_amp: bool = True,
        gradient_checkpointing: bool = True,
        reserve_gb: float = 2.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
        # Derived (RTX 4090 default; override per profile):
        self.total_vram_gb = 24.0
        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 — use load_all_configs in production)."""
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}hardware.gin")
    return HardwareConfig()
```

### 3.3 `src/conf/models_conf.py`

```python
from __future__ import annotations

import gin


@gin.configurable
class ModelsConfig:
    """Model checkpoints + architecture switches.

    Everything that defines WHAT model is built: backbone choice, paths to
    pretrained weights, MONA placement, gate init, asym vs shared encoder.
    Changes here mean a new architecture experiment.
    """

    def __init__(
        self,
        # Checkpoints.
        dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
        dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
        lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
        stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
        # Backbone selection.
        backbone: str = "dinov3",
        shared_encoder: bool = True,
        baseline_mode: bool = False,
        # Fusion.
        init_gate: float = 0.7,
        # MONA (DINOv3).
        mona_bottleneck: int = 64,
        mona_last_n_blocks: int = 12,
        # StripNet-specific.
        stripnet_freeze: bool = True,
        stripnet_mona_last_n_stages: int = 2,
    ) -> None:
        self.dino_web_path = dino_web_path
        self.dino_sat_path = dino_sat_path
        self.lrsclip_path = lrsclip_path
        self.stripnet_path = stripnet_path
        self.backbone = backbone
        self.shared_encoder = shared_encoder
        self.baseline_mode = baseline_mode
        self.init_gate = init_gate
        self.mona_bottleneck = mona_bottleneck
        self.mona_last_n_blocks = mona_last_n_blocks
        self.stripnet_freeze = stripnet_freeze
        self.stripnet_mona_last_n_stages = stripnet_mona_last_n_stages


def get_models_cfg(path2cfg: str) -> ModelsConfig:
    """Load ONLY models config (TESTING ONLY — use load_all_configs in production)."""
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}models.gin")
    return ModelsConfig()
```

### 3.4 `src/conf/training_conf.py`

```python
from __future__ import annotations

import gin


@gin.configurable
class TrainingConfig:
    """Training recipe: loss + optimizer + sampler.

    These three move together when you tune learning. Changing tau usually
    pairs with changing lr; switching sampler_type usually pairs with
    re-tuning loss weights. Keeping them in one config matches the actual
    workflow of running ablations.
    """

    def __init__(
        self,
        # --- Loss ---
        loss_type: str = "symmetric",
        tau_init: float = 0.07,
        tau_min: float = 0.01,
        tau_max: float = 0.1,
        learnable_temperature: bool = True,
        label_smoothing: float = 0.1,
        weight_q2g: float = 0.6,
        weight_g2q: float = 0.4,
        hard_mining_k: int = 0,
        neg_bank_size: int = 0,
        # --- Optimizer ---
        learning_rate: float = 1e-4,
        text_lr_factor: float = 0.1,
        stripnet_backbone_lr_factor: float = 0.1,
        weight_decay: float = 1e-4,
        grad_clip: float = 1.0,
        # --- Sampler ---
        sampler_type: str = "mutex",
        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,
    ) -> None:
        # Loss.
        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
        self.weight_q2g = weight_q2g
        self.weight_g2q = weight_g2q
        self.hard_mining_k = hard_mining_k
        self.neg_bank_size = neg_bank_size
        # Optimizer.
        self.learning_rate = learning_rate
        self.text_lr_factor = text_lr_factor
        self.stripnet_backbone_lr_factor = stripnet_backbone_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


def get_training_cfg(path2cfg: str) -> TrainingConfig:
    """Load ONLY training config (TESTING ONLY — use load_all_configs in production)."""
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}training.gin")
    return TrainingConfig()
```

### 3.5 `src/conf/tracking_conf.py`

```python
from __future__ import annotations

import gin


@gin.configurable
class TrackingConfig:
    """Experiment tracking + diagnostics.

    Independent axis: changing these flags does not affect training results,
    only what is observed/recorded.
    """

    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 — use load_all_configs in production)."""
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}tracking.gin")
    return TrackingConfig()
```

### 3.6 `src/conf/config_loader.py` — продакшен-вход

```python
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_conf import ModelsConfig
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__)


def load_all_configs(path2cfg: str) -> dict[str, Any]:
    """Parse ALL .gin files in path2cfg and return all config objects.

    This is the PRODUCTION entry point — main() calls this once. Individual
    get_*_cfg() loaders exist only for unit tests / notebooks.

    Args:
        path2cfg: Path to config directory (WITH trailing slash).

    Returns:
        Dict with config objects keyed by name:
            {
                "pipeline": PipelineConfig,
                "hardware": HardwareConfig,
                "models":   ModelsConfig,
                "training": TrainingConfig,
                "tracking": TrackingConfig,
            }

    Raises:
        FileNotFoundError: If path2cfg contains no .gin files.
    """
    cfg_dir = Path(path2cfg)
    gin_files = sorted(cfg_dir.glob("*.gin"))
    if not gin_files:
        raise FileNotFoundError(f"No .gin files found in {cfg_dir}")

    # MANDATORY: reset gin global state before parsing — without clear_config(),
    # parameters from previous parses accumulate (gin holds global bindings).
    gin.clear_config()
    gin.parse_config_files_and_bindings(
        config_files=[str(f) for f in gin_files],
        bindings=[],
    )
    logger.info("Loaded %d gin files from %s", len(gin_files), cfg_dir)

    # Instantiate AFTER all bindings are parsed.
    return {
        "pipeline": PipelineConfig(),
        "hardware": HardwareConfig(),
        "models":   ModelsConfig(),
        "training": TrainingConfig(),
        "tracking": TrackingConfig(),
    }
```

---

## 4. `InfoNCELoss`: убрать двойную gin-регистрацию

**Текущая проблема.** В `src/losses/multi_infonce.py` `InfoNCELoss` декорирован `@gin.configurable`, и в `gtauav_balanced.gin` есть **обе** группы биндингов:

```gin
TrainConfigGTAUAV.tau_init = 0.07          # <-- читается мега-конфигом
TrainConfigGTAUAV.weight_q2g = 0.6
# ...
InfoNCELoss.temperature_init = 0.07        # <-- читается самой InfoNCELoss
InfoNCELoss.weight_q2g = 0.6
# ...
```

Это два источника правды на одни и те же значения. Если кто-то поменяет одну строку и забудет другую — обучение пойдёт с рассинхроном между логированием и реальным `tau`. Тихая ошибка, которую очень тяжело отладить.

**Решение: один источник правды — `TrainingConfig`.** `InfoNCELoss` теряет `@gin.configurable` и принимает параметры обычными аргументами:

```python
# src/losses/multi_infonce.py
class InfoNCELoss(nn.Module):
    """Symmetric InfoNCE with learnable temperature.

    Note: NOT @gin.configurable. All parameters come from TrainingConfig
    via explicit kwargs, which keeps a single source of truth.
    """

    def __init__(
        self,
        temperature_init: float = 0.07,
        temperature_min: float = 0.01,
        temperature_max: float = 0.1,
        learnable_temperature: bool = True,
        label_smoothing: float = 0.1,
        weight_q2g: float = 0.6,
        weight_g2q: float = 0.4,
        hard_mining_k: int = 0,
    ) -> None:
        super().__init__()
        # ... (existing implementation)
```

В `Trainer`:

```python
loss_fn = InfoNCELoss(
    temperature_init=training_cfg.tau_init,
    temperature_min=training_cfg.tau_min,
    temperature_max=training_cfg.tau_max,
    learnable_temperature=training_cfg.learnable_temperature,
    label_smoothing=training_cfg.label_smoothing,
    weight_q2g=training_cfg.weight_q2g,
    weight_g2q=training_cfg.weight_g2q,
    hard_mining_k=training_cfg.hard_mining_k,
)
```

В `training.gin` остаются **только** `TrainingConfig.*` биндинги. Все `InfoNCELoss.*` строки **удаляются** из всех gin-файлов.

**Принцип общий:** `@gin.configurable` идёт **только** на классы из `src/conf/`. Никакие классы из `src/models/`, `src/losses/`, `src/datasets/` не должны быть gin-configurable. Они принимают параметры явно через `__init__` и собираются `Trainer`-ом из объектов конфига.

---

## 5. Точка входа — `src/main.py`

```python
from __future__ import annotations

import logging

import coloredlogs

from src.conf.config_loader import load_all_configs
from src.training.trainer import Trainer
from src.utils.path_utils import get_proj_dir
from src.utils.seed_utils import set_seed

logger = logging.getLogger("caption_test")


def main() -> None:
    """Entry point: load all configs from in/config_files/ and run training.

    No argparse. No CLI flags. No env vars. Every parameter lives in .gin
    files under in/config_files/. To switch experiments, copy a different
    preset over in/config_files/ before running.
    """
    coloredlogs.install(
        level="INFO",
        logger=logger,
        fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
    )

    proj_dir = get_proj_dir()
    path2cfg = f"{proj_dir}in/config_files/"

    # ONE call loads everything; gin.clear_config() runs inside.
    configs = load_all_configs(path2cfg)

    # Reproducibility — single point of seed setting.
    set_seed(configs["pipeline"].seed)

    # Pass configs explicitly — Trainer never queries gin global state.
    trainer = Trainer(
        pipeline_cfg=configs["pipeline"],
        hardware_cfg=configs["hardware"],
        models_cfg=configs["models"],
        training_cfg=configs["training"],
        tracking_cfg=configs["tracking"],
    )
    trainer.run()


if __name__ == "__main__":
    main()
```

### `src/utils/path_utils.py`

```python
from __future__ import annotations

from pathlib import Path

# Markers that identify the project root.
_MARKERS: tuple[str, ...] = ("pyproject.toml", ".git", "in")


def get_proj_dir() -> str:
    """Return absolute project root directory with trailing slash.

    Walks up from this file's directory until it finds one of the markers
    (pyproject.toml, .git, or `in/`). Falls back with a clear error if not
    found within 10 levels.

    Returns:
        Project root path as string 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`

```python
from __future__ import annotations

import random

import numpy as np
import torch


def set_seed(seed: int = 42) -> None:
    """Fix all RNG seeds for reproducibility.

    Args:
        seed: Integer seed applied to Python random, NumPy, and PyTorch
            (CPU + all CUDA devices).
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
```

### `src/utils/io_utils.py`

```python
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 (including KeyboardInterrupt / SIGTERM), the temp file
    is removed. This makes --resume safe: a partial checkpoint never ends
    up at the destination path.

    Args:
        obj: Anything torch.save can handle (state dict, full model, etc.).
        path: Destination path. Parent directory is created if missing.

    Raises:
        Re-raises any error from torch.save after cleaning up the temp file.
    """
    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.

    Call before starting a new training stage or when switching models.
    """
    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)
```

---

## 6. Пример gin-пресета `gtauav_balanced/`

### `presets/gtauav_balanced/pipeline.gin`

```gin
# What to train on, where to save, how long.
PipelineConfig.train_json = 'meta/train_80.json'
PipelineConfig.test_json = 'meta/test_20.json'
PipelineConfig.rgb_root = '/data/GTA-UAV-LR'
PipelineConfig.caption_root = '/data/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
```

### `presets/gtauav_balanced/hardware.gin`

```gin
# RTX 4090 profile.
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
```

### `presets/gtauav_balanced/models.gin`

```gin
# Architecture: shared DINOv3 WEB encoder, MONA in last 12 of 24 blocks, with text.
ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
ModelsConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
ModelsConfig.backbone = 'dinov3'
ModelsConfig.shared_encoder = True
ModelsConfig.baseline_mode = False
ModelsConfig.init_gate = 0.7
ModelsConfig.mona_bottleneck = 64
ModelsConfig.mona_last_n_blocks = 12
```

### `presets/gtauav_balanced/training.gin`

```gin
# Loss + optimizer + sampler — the training recipe.
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.weight_q2g = 0.6
TrainingConfig.weight_g2q = 0.4
TrainingConfig.hard_mining_k = 0
TrainingConfig.neg_bank_size = 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
```

### `presets/gtauav_balanced/tracking.gin`

```gin
# Diagnostics off by default; flip on per experiment.
TrackingConfig.use_wandb = False
TrackingConfig.use_tb = True
TrackingConfig.log_grad_norms = True
TrackingConfig.use_gradcam = False
TrackingConfig.gradcam_every = 5
TrackingConfig.use_profiler = False
```

> **Важно:** ни в одном `.gin` файле нет ни `TrainConfigGTAUAV.*`, ни `InfoNCELoss.*` биндингов. Только 5 классов из `src/conf/`. Каждый параметр имеет ровно одно место, где он определён.

---

## 7. Декомпозиция `train_gtauav.py` (1296 → ~5 файлов по 100–250 строк)

| Сейчас в `train_gtauav.py` | Куда переносится | Размер |
|---|---|---|
| `TrainConfigGTAUAV` (`@dataclass + @gin.configurable`) | **Удаляется**, заменяется 5 классами в `src/conf/` | — |
| Module-level constants `_RGB_ROOT`, `_DINO_WEB`, и т.д. | **Удаляются**, дефолты → `__init__` конфиг-классов | — |
| `_set_seed()` | `src/utils/seed_utils.py::set_seed` | ~10 строк |
| `_atomic_save()` | `src/utils/io_utils.py::atomic_save_torch` | ~15 строк |
| `_clear_vram()` | `src/utils/io_utils.py::clear_vram` | ~10 строк |
| `_build_param_groups()` | `src/training/trainer.py::Trainer._build_param_groups` (метод) | ~30 строк |
| `_cosine_warmup_schedule()` | `src/training/trainer.py` (модульная функция) | ~10 строк |
| `_embed_drone_queries()` (`@torch.no_grad`) | `src/training/trainer.py::Trainer._embed_drone_queries` (`@torch.inference_mode`) | ~30 строк |
| `_evaluate()` (`@torch.no_grad`) | `src/eval/evaluator.py::evaluate` (`@torch.inference_mode`) | ~150 строк |
| `CSVLogger` | `src/training/csv_logger.py` | ~80 строк |
| `train()` (700+ строк) | `src/training/trainer.py::Trainer.run` + приватные методы | ~250 строк |
| `main()` (с argparse) | `src/main.py` (без argparse) | ~30 строк |

### Скелет `Trainer`

```python
# src/training/trainer.py
from __future__ import annotations

import logging
from pathlib import Path

# ... imports

logger = logging.getLogger(__name__)


class Trainer:
    """Orchestrates the full GTA-UAV training pipeline.

    Not @gin.configurable — all parameters arrive as config objects.
    """

    def __init__(
        self,
        pipeline_cfg: PipelineConfig,
        hardware_cfg: HardwareConfig,
        models_cfg: ModelsConfig,
        training_cfg: TrainingConfig,
        tracking_cfg: TrackingConfig,
    ) -> None:
        self.pipeline_cfg = pipeline_cfg
        self.hardware_cfg = hardware_cfg
        self.models_cfg = models_cfg
        self.training_cfg = training_cfg
        self.tracking_cfg = tracking_cfg

        # Populated lazily in _setup / _build_*.
        self.output_dir: Path | None = None
        self.model: AsymmetricEncoder | None = None
        self.loss_fn: nn.Module | None = None
        self.optimizer: torch.optim.Optimizer | None = None
        self.scheduler: torch.optim.lr_scheduler.LambdaLR | None = None
        self.train_loader: DataLoader | None = None
        self.test_loader: DataLoader | None = None
        self.tracker: ExperimentTracker | None = None
        self.csv_logger: CSVLogger | None = None

    def run(self) -> None:
        """Full pipeline: setup → train → eval → save."""
        clear_vram()
        self._setup_output_dir()
        self._build_tracker()
        self._build_model()
        self._build_loss()
        self._build_data_loaders()
        self._build_optimizer_and_scheduler()
        try:
            self._train_loop()
            self._final_evaluation()
        finally:
            self._cleanup()

    # --- Private helpers below (one method per concern) ---

    def _setup_output_dir(self) -> None: ...
    def _build_tracker(self) -> None: ...
    def _build_model(self) -> None: ...
    def _build_loss(self) -> None: ...
    def _build_data_loaders(self) -> None: ...
    def _build_optimizer_and_scheduler(self) -> None: ...
    def _build_param_groups(self) -> list[dict]: ...
    def _train_loop(self) -> None: ...
    def _train_one_epoch(self, epoch: int) -> dict[str, float]: ...
    def _evaluate_epoch(self, epoch: int) -> None: ...
    def _save_checkpoint(self, epoch: int) -> None: ...
    def _final_evaluation(self) -> None: ...
    def _cleanup(self) -> None: ...
```

Каждый приватный метод — 20–60 строк, одна ответственность. Читать и тестировать на порядок легче, чем 700-строчную функцию.

---

## 8. `@torch.inference_mode()` вместо `@torch.no_grad()`

В текущем коде:

```python
# src/training/train_gtauav.py
@torch.no_grad()
def _embed_drone_queries(...): ...

@torch.no_grad()
def _evaluate(...): ...
```

Стандарт §4.4 требует `@torch.inference_mode()`. Это не просто «нагляднее»: `inference_mode()` дополнительно отключает version counter на тензорах, что даёт небольшой speedup и блокирует случайные in-place правки, которые при `no_grad()` тихо пройдут.

При переносе в `src/eval/evaluator.py` и `src/training/trainer.py`:

```python
# src/eval/evaluator.py
@torch.inference_mode()
def evaluate(
    model: nn.Module,
    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."""
    # ... rest unchanged ...
```

---

## 9. Чеклист по review (полный, по стандарту §6)

| Пункт | Текущее состояние | После рефакторинга |
|---|---|---|
| `from __future__ import annotations` первой строкой | ⚠️ `train_gtauav.py` ✓; остальные требуют аудита | ✅ во всех новых файлах |
| Все функции/методы имеют type hints | ⚠️ частично (`_atomic_save(obj: dict, ...)` слабо типизировано) | ✅ строгие type hints |
| Google-style docstrings на публичных классах/функциях | ✅ в основном есть | ✅ + покрытие 100% |
| `@gin.configurable` только на классах | ❌ на `dataclass` (запрещено) | ✅ только на классах из `src/conf/` |
| Нет `dataclass` + gin | ❌ `TrainConfigGTAUAV` нарушает | ✅ удалён |
| Нет `argparse` | ❌ 15+ CLI флагов в `main()` | ✅ убран полностью |
| Нет захардкоженных model ID / промптов / размеров | ❌ `_DINO_WEB`, `_DINO_SAT`, `_LRSCLIP`, `_RGB_ROOT` на module-level | ✅ только дефолты в `__init__` конфигов |
| `gin.clear_config()` перед каждой загрузкой | ❌ не вызывается нигде | ✅ внутри `load_all_configs()` |
| Один источник правды для каждого параметра | ❌ дубль `TrainConfigGTAUAV.tau_init` ↔ `InfoNCELoss.temperature_init` | ✅ только `TrainingConfig.*` |
| Модели выгружаются после использования | ⚠️ `_clear_vram` в начале, нет `del model` в конце | ✅ `_cleanup()` в `Trainer` |
| Файлы сохраняются атомарно (temp + replace + cleanup на ошибке) | ⚠️ нет `try/except` для очистки `.tmp` | ✅ `atomic_save_torch` с `tempfile.mkstemp` |
| Seed установлен | ✅ `_set_seed(42)` | ✅ через `src/utils/seed_utils.py::set_seed` |
| `@torch.inference_mode()` на inference-функциях | ❌ используется `@torch.no_grad()` | ✅ заменено везде |
| Английский язык кода/комментариев | ✅ хорошо | ✅ |
| Импорты: stdlib → third-party → local, разделены пустыми строками | ⚠️ требует аудита | ✅ во всех новых файлах |

---

## 10. План миграции — 7 коммитов

> Каждый коммит самодостаточен и оставляет код рабочим. После каждого коммита прогон 1 эпохи на маленьком сабсете и сверка `r@1_q2g`/`loss` до 4-го знака.

| # | Коммит | Что делается | Что не ломается |
|---|---|---|---|
| 1 | **utils** | Создать `src/utils/{io_utils,seed_utils,path_utils}.py`. Переключить `train_gtauav.py` на новые имена | Старый код продолжает работать |
| 2 | **conf infrastructure** | Создать 5 классов в `src/conf/` + `config_loader.py`. **Не использовать.** | Существующий `TrainConfigGTAUAV` остаётся |
| 3 | **evaluator** | Вынести `_evaluate` → `src/eval/evaluator.py`, `@torch.inference_mode()` | `train_gtauav.py` импортирует оттуда |
| 4 | **csv logger + trainer skeleton** | Вынести `CSVLogger` → `src/training/csv_logger.py`. Создать `Trainer` (пока пустой), но не использовать | Параллельно работают оба пути |
| 5 | **Trainer.run() реализация** | Перенести логику `train()` в `Trainer.run()` методы. Создать первый пресет `presets/gtauav_balanced/`. Создать `src/main.py` | Обе точки входа работают; сравнить метрики |
| 6 | **InfoNCELoss decouple** | Убрать `@gin.configurable` с `InfoNCELoss`. Удалить `InfoNCELoss.*` биндинги из всех `.gin` | `TrainingConfig` единственный источник |
| 7 | **cleanup** | Удалить `TrainConfigGTAUAV`, `argparse`, `train_gtauav.py::main`, старую `conf/` | Только новый путь |

**Контрольная точка**: после коммита 5 запустить **обе** точки входа (`python -m src.training.train_gtauav --config conf/gtauav_balanced.gin` и `python -m src.main` с пресетом `gtauav_balanced`) на 1 эпохе на 16 батчах. Метрики `r@1_q2g`, `r@5_q2g`, `loss`, `tau`, `gate_q`, `gate_g` должны совпадать до 4-го знака. Любое расхождение — баг рефакторинга, не логики.

---

## 11. Что **не** трогать

Эти модули содержат рабочую ML-логику и не подлежат рефакторингу в рамках этой задачи:

- `src/models/asymmetric_encoder.py`
- `src/models/dgtrs/`
- `src/models/adapters.py` (MONA, LoRA)
- `src/datasets/gtauav_dataset.py`
- `src/datasets/dynamic_similarity_sampler.py`
- `src/datasets/mutually_exclusive_sampler.py`
- `src/datasets/embedding_cache.py`
- `src/losses/multi_infonce.py` — **только** убрать `@gin.configurable`, остальное не трогать
- `src/losses/weighted_infonce.py`, `src/losses/hard_negatives.py`
- `src/training/grad_monitor.py`, `gradcam.py`, `profiling.py`, `plot_metrics.py`, `trackers.py`

Цель этого рефакторинга — **только конфиг и точка входа**. Никаких изменений в формуле обучения, в модельной архитектуре, в семплерах или в формуле loss. Метрики после рефакторинга должны быть идентичны метрикам до — иначе это регрессия, а не рефакторинг.

---

## 12. Опционально, но окупится: smoke-test

```python
# tests/test_smoke.py
from __future__ import annotations

from pathlib import Path

import pytest

from src.conf.config_loader import load_all_configs


def test_load_all_configs_returns_5_keys(tmp_path: Path) -> None:
    """load_all_configs returns exactly the 5 expected keys."""
    # Write minimal .gin files into tmp_path/.
    (tmp_path / "pipeline.gin").write_text("PipelineConfig.epochs = 1\n")
    (tmp_path / "hardware.gin").write_text("HardwareConfig.batch_size = 2\n")
    (tmp_path / "models.gin").write_text("ModelsConfig.backbone = 'dinov3'\n")
    (tmp_path / "training.gin").write_text("TrainingConfig.tau_init = 0.07\n")
    (tmp_path / "tracking.gin").write_text("TrackingConfig.use_tb = False\n")

    cfgs = load_all_configs(str(tmp_path) + "/")
    assert set(cfgs.keys()) == {"pipeline", "hardware", "models", "training", "tracking"}
    assert cfgs["pipeline"].epochs == 1
    assert cfgs["hardware"].batch_size == 2


def test_load_all_configs_clears_state_between_calls(tmp_path: Path) -> None:
    """Two calls with different .gin do not leak state."""
    # Call 1: epochs=10
    (tmp_path / "pipeline.gin").write_text("PipelineConfig.epochs = 10\n")
    for name in ("hardware", "models", "training", "tracking"):
        (tmp_path / f"{name}.gin").write_text("")  # empty is fine
    cfgs1 = load_all_configs(str(tmp_path) + "/")
    assert cfgs1["pipeline"].epochs == 10

    # Call 2: epochs=20 — must NOT inherit 10 from prior call.
    (tmp_path / "pipeline.gin").write_text("PipelineConfig.epochs = 20\n")
    cfgs2 = load_all_configs(str(tmp_path) + "/")
    assert cfgs2["pipeline"].epochs == 20  # would fail if clear_config absent
```

Второй тест ловит самый коварный баг при работе с gin — пропущенный `clear_config()`. Один раз отладишь — ловится за полсекунды на каждом запуске CI.

---

## Резюме одной строкой

**Удалить `TrainConfigGTAUAV` (`@dataclass + @gin.configurable`) и argparse → разделить на 5 узких конфиг-классов в `src/conf/` → грузить через единственный `load_all_configs()` с `gin.clear_config()` → разрезать `train_gtauav.py` (1296 строк) на `Trainer` + `evaluator` + `csv_logger` + `utils` + `main` → убрать дубль gin-биндингов между `TrainConfigGTAUAV` и `InfoNCELoss` (один источник правды — `TrainingConfig`) → `@torch.no_grad` → `@torch.inference_mode`.**