From 5c48b6c8fdcf7486d6184c9e8bb92aa35be55c67 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Wed, 29 Apr 2026 11:55:38 +0300 Subject: [PATCH] claude_refactor_v2: temp changes before src changes --- in/config_files/hardware.gin | 1 + in/config_files/models.gin | 1 + in/config_files/pipeline.gin | 1 + in/config_files/tracking.gin | 1 + in/config_files/training.gin | 1 + presets/gtauav_balanced/models.gin | 11 + presets/gtauav_balanced/pipeline.gin | 13 + refactor_proposal_v2.md | 1089 ++++++++++++++++++++++++++ src/conf/config_loader.py | 62 ++ src/conf/hardware_conf.py | 43 + src/conf/models_conf.py | 54 ++ src/conf/pipeline_conf.py | 57 ++ src/conf/tracking_conf.py | 48 ++ src/conf/training_conf.py | 79 ++ 14 files changed, 1461 insertions(+) create mode 100644 in/config_files/hardware.gin create mode 100644 in/config_files/models.gin create mode 100644 in/config_files/pipeline.gin create mode 100644 in/config_files/tracking.gin create mode 100644 in/config_files/training.gin create mode 100644 presets/gtauav_balanced/models.gin create mode 100644 presets/gtauav_balanced/pipeline.gin create mode 100644 refactor_proposal_v2.md create mode 100644 src/conf/config_loader.py create mode 100644 src/conf/hardware_conf.py create mode 100644 src/conf/models_conf.py create mode 100644 src/conf/pipeline_conf.py create mode 100644 src/conf/tracking_conf.py create mode 100644 src/conf/training_conf.py diff --git a/in/config_files/hardware.gin b/in/config_files/hardware.gin new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/in/config_files/hardware.gin @@ -0,0 +1 @@ + diff --git a/in/config_files/models.gin b/in/config_files/models.gin new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/in/config_files/models.gin @@ -0,0 +1 @@ + diff --git a/in/config_files/pipeline.gin b/in/config_files/pipeline.gin new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/in/config_files/pipeline.gin @@ -0,0 +1 @@ + diff --git a/in/config_files/tracking.gin b/in/config_files/tracking.gin new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/in/config_files/tracking.gin @@ -0,0 +1 @@ + diff --git a/in/config_files/training.gin b/in/config_files/training.gin new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/in/config_files/training.gin @@ -0,0 +1 @@ + diff --git a/presets/gtauav_balanced/models.gin b/presets/gtauav_balanced/models.gin new file mode 100644 index 0000000..7c7f289 --- /dev/null +++ b/presets/gtauav_balanced/models.gin @@ -0,0 +1,11 @@ +# Architecture: shared DINOv3 WEB encoder + DGTRS-CLIP text + MONA in last 12/24 blocks. +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.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth' +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 \ No newline at end of file diff --git a/presets/gtauav_balanced/pipeline.gin b/presets/gtauav_balanced/pipeline.gin new file mode 100644 index 0000000..20096fd --- /dev/null +++ b/presets/gtauav_balanced/pipeline.gin @@ -0,0 +1,13 @@ +# 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 = '/home/servml/Документы/datasets/GTA-UAV-LR' +PipelineConfig.caption_root = '/home/servml/Документы/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 + diff --git a/refactor_proposal_v2.md b/refactor_proposal_v2.md new file mode 100644 index 0000000..5cd1522 --- /dev/null +++ b/refactor_proposal_v2.md @@ -0,0 +1,1089 @@ +# Рефакторинг `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`.** diff --git a/src/conf/config_loader.py b/src/conf/config_loader.py new file mode 100644 index 0000000..0125d74 --- /dev/null +++ b/src/conf/config_loader.py @@ -0,0 +1,62 @@ +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(), + } + diff --git a/src/conf/hardware_conf.py b/src/conf/hardware_conf.py new file mode 100644 index 0000000..0d0e7b9 --- /dev/null +++ b/src/conf/hardware_conf.py @@ -0,0 +1,43 @@ +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() + diff --git a/src/conf/models_conf.py b/src/conf/models_conf.py new file mode 100644 index 0000000..448ac0b --- /dev/null +++ b/src/conf/models_conf.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import gin + + +@gin.configurable +class ModelsConfig: + """Model checkpoints + architecture switches. + + Default checkpoint paths are relative to the project root (matching the + repository layout: nn_models/DINO_WEB/, nn_models/DINO_SAT/, etc.). + These are gitignored and must be downloaded separately — see README. + """ + + def __init__( + self, + # Checkpoints — relative to project root, defaults match repo layout. + 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() + diff --git a/src/conf/pipeline_conf.py b/src/conf/pipeline_conf.py new file mode 100644 index 0000000..a2f41ea --- /dev/null +++ b/src/conf/pipeline_conf.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import gin + + +@gin.configurable +class PipelineConfig: + """Pipeline orchestration: data IO, training schedule, output, resume. + + Defaults match the current `belka_refactor` HEAD: hardcoded servml paths + are preserved verbatim. To switch to a different machine, override in + pipeline.gin — never edit defaults here. + """ + + def __init__( + self, + # Data paths (defaults match servml workstation). + train_json: str = "meta/train_80.json", + test_json: str = "meta/test_20.json", + rgb_root: str = "/home/servml/Документы/datasets/GTA-UAV-LR", + caption_root: str = "/home/servml/Документы/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 — 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() + diff --git a/src/conf/tracking_conf.py b/src/conf/tracking_conf.py new file mode 100644 index 0000000..d4974a6 --- /dev/null +++ b/src/conf/tracking_conf.py @@ -0,0 +1,48 @@ +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() + diff --git a/src/conf/training_conf.py b/src/conf/training_conf.py new file mode 100644 index 0000000..264c450 --- /dev/null +++ b/src/conf/training_conf.py @@ -0,0 +1,79 @@ +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() + +