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 Критические (запрещены прямо)

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

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

# 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 смены пресета:

# Вместо 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

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

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

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

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

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 — продакшен-вход

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 есть обе группы биндингов:

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 и принимает параметры обычными аргументами:

# 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:

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

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

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

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

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

# 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

# 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

# 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

# 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

# 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

# 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()

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

# 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:

# 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

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