From 6b928634f463aee6232c58222c360c975f6f2b9b Mon Sep 17 00:00:00 2001 From: pikaliov Date: Wed, 29 Apr 2026 14:43:42 +0300 Subject: [PATCH] clean up --- refactor_proposal_v2.md | 1089 --------------------------------------- 1 file changed, 1089 deletions(-) delete mode 100644 refactor_proposal_v2.md diff --git a/refactor_proposal_v2.md b/refactor_proposal_v2.md deleted file mode 100644 index 5cd1522..0000000 --- a/refactor_proposal_v2.md +++ /dev/null @@ -1,1089 +0,0 @@ -# Рефакторинг `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`.**