Pikaliov/caption-test

Fork 0

Files

pikaliov 89cb8ab0f7 claude_refactor_v3: New train_gtauav.py, added entry point main.py, added utils

2026-05-04 11:20:14 +03:00

16 KiB

Raw Blame History

Шаг 4а — Что изменилось

1. Новая точка входа — `src/main.py`

Запуск тренировки переехал из src/training/train_gtauav.py::main() в отдельный модуль src/main.py.

Старый запуск:

python src/training/train_gtauav.py --config conf/gtauav_balanced.gin

Новый запуск:

python -m src.main gtauav_balanced

src/main.py:

Читает имя пресета из sys.argv[1] (один позиционный аргумент)
Резолвит корень проекта через get_proj_dir() (поиск по маркерам pyproject.toml/.git/in/)
Формирует path2cfg = f"{proj_dir}in/config_files/" буквально по REQUIREMENTS_GIN_STYLE.md §5
Вызывает load_all_configs(path2cfg, preset_name) — двухпроходная загрузка из _common-файлов и пресет-директории
Передаёт 6 объектов конфига в train(...) именованными аргументами

Никакого argparse, никаких CLI-overrides — все параметры в .gin-файлах.

2. Изменённый `src/training/train_gtauav.py`

2.1 — Удалено

import argparse, import gin, from dataclasses import dataclass, field
Класс TrainConfigGTAUAV (@dataclass + @gin.configurable) — все его поля переехали в 6 классов в src/conf/
Module-level константы _RGB_ROOT, _CAPTION_ROOT, _TRAIN_JSON, _TEST_JSON, _DINO_WEB, _DINO_SAT, _LRSCLIP
Функция main() с argparse и CLI-overrides

2.2 — Изменена сигнатура `train()`

Было:

def train(cfg: TrainConfigGTAUAV) -> None:

Стало:

def train(
    pipeline_cfg: PipelineConfig,
    hardware_cfg: HardwareConfig,
    training_cfg: TrainingConfig,
    tracking_cfg: TrackingConfig,
    models_common_cfg: ModelsCommonConfig,
    models_cfg: DINOv3ModelsConfig | StripNetModelsConfig | SOFIAv1ModelsConfig | SOFIAv71ModelsConfig,
) -> None:

2.3 — Обращения `cfg.xxx` переписаны

По карте уникальных полей:

cfg.train_json, cfg.rgb_root, cfg.epochs, cfg.output_dir, cfg.seed, ... → pipeline_cfg.*
cfg.batch_size, cfg.grad_accum_steps, cfg.use_amp, cfg.gradient_checkpointing, ... → hardware_cfg.*
cfg.tau_init, cfg.learning_rate, cfg.sampler_type, cfg.dss_*, ... → training_cfg.*
cfg.use_wandb, cfg.use_tb, cfg.use_gradcam, cfg.use_profiler, ... → tracking_cfg.*
cfg.backbone, cfg.baseline_mode, cfg.init_gate, cfg.lrsclip_path → models_common_cfg.*
cfg.dino_web_path, cfg.shared_encoder, cfg.mona_* (DINOv3-only) → models_cfg.*
cfg.stripnet_* (StripNet-only) → models_cfg.*
cfg.sofia_preset → models_cfg.variant_label, cfg.sofia_d_descriptor → models_cfg.d_descriptor, cfg.sofia_use_text_film_*, cfg.sofia_mamba_* → models_cfg.*
cfg.sofia_v1_variant → models_cfg.variant_label, cfg.sofia_v1_* → models_cfg.*

2.4 — Sofia-модели строятся напрямую из gin

Раньше Sofia v7.1 строился через preset-фабрику + точечные overrides:

# было
preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}
sofia_cfg = preset_map[cfg.sofia_preset]()      # строит SOFIAConfig с дефолтами размера
sofia_cfg.d_descriptor = cfg.sofia_d_descriptor # потом 8 overrides
sofia_cfg.use_text_film_uav = ...
...

Теперь SOFIAConfig(...) собирается напрямую из всех 40+ полей SOFIAv71ModelsConfig:

# стало
sofia_cfg = SOFIAConfig(
    input_size=models_cfg.input_size,
    embed_dims=list(models_cfg.embed_dims),       # все 4 dims из gin
    depths=list(models_cfg.depths),
    mamba_extra_kwargs=dict(models_cfg.mamba_extra_kwargs),
    ...                                            # и все остальные 35+ полей
)

Преимущество: каждый размер Sofia (Tiny/M/L) — это отдельный presets/<name>/models.gin со всеми полями явно. Не нужно знать, что кладёт sofia_tiny_config() в дефолтах. Один источник правды — gin.

Аналогично для Sofia v1: SOFIAv1Config(...) строится из полей SOFIAv1ModelsConfig.

2.5 — Direct execution убран

if __name__ == "__main__":
    raise SystemExit(
        "Direct execution removed. Use: python -m src.main <preset_name>",
    )

3. Проверка передаваемых путей в энкодерах и бэкбонах

После того как TrainConfigGTAUAV исчез, поля стали раскиданы по 4 семейным Models*Config-классам. Для того чтобы поведение точно совпадало со старым кодом, в каждой ветке сборки модели мы передаём те же значения, что приходили раньше из cfg.* — даже если поле теперь не имеет смысла для активного бэкбона.

3.1 — Ветка `sofia_v71`

SOFIAFusionEncoder(
    sofia_cfg=...,                                     # из SOFIAv71ModelsConfig
    lrsclip_path=models_common_cfg.lrsclip_path,       # общий путь к DGTRS-CLIP
    init_gate=models_common_cfg.init_gate,
    baseline_mode=models_common_cfg.baseline_mode,
    lora_rank=models_cfg.lora_rank,
    device=hardware_cfg.device,
)

— ничего лишнего, всё из gin.

3.2 — Ветка `sofia_v1`

SOFIAv1FusionEncoder(
    sofia_cfg=SOFIAv1Config(variant=models_cfg.variant_label, ...),
    lrsclip_path=models_common_cfg.lrsclip_path,
    init_gate=models_common_cfg.init_gate,
    baseline_mode=models_common_cfg.baseline_mode,
    lora_rank=models_cfg.lora_rank,
    device=hardware_cfg.device,
)

— симметрично с v7.1.

3.3 — Ветка `stripnet`

AsymmetricEncoder(
    dino_web_path="nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",  # старый _DINO_WEB
    dino_sat_path="nn_models/DINO_SAT/model.safetensors",                     # старый _DINO_SAT
    lrsclip_path=models_common_cfg.lrsclip_path,
    init_gate=models_common_cfg.init_gate,
    baseline_mode=models_common_cfg.baseline_mode,
    shared_encoder=True,                                                       # для StripNet всегда True
    mona_bottleneck=64,                                                        # старый дефолт TrainConfigGTAUAV
    mona_last_n_blocks=12,                                                     # старый дефолт
    device=hardware_cfg.device,
    backbone=backbone,
    stripnet_path=models_cfg.stripnet_path,
    stripnet_mona_last_n_stages=models_cfg.stripnet_mona_last_n_stages,
    stripnet_freeze=models_cfg.stripnet_freeze,
)

Почему DINO-пути передаются для StripNet: AsymmetricEncoder.__init__ принимает все 13 параметров независимо от backbone. Для StripNet-режима DINO-пути игнорируются (модель строит StripNetEncoder, не DINOv3ViT). Старый код передавал те же _DINO_WEB/_DINO_SAT всегда — мы воспроизводим точно. Семантика одинакова.

Почему shared_encoder=True: внутри AsymmetricEncoder.__init__ на строке if backbone == "stripnet": self.shared_encoder = True — значение всё равно перезаписывается. Передаём True для семантической чистоты.

Почему mona_bottleneck=64/mona_last_n_blocks=12: mona_bottleneck=64 используется при inject_conv_mona_into_stripnet(...) для StripNet — нужно валидное значение. Старый код всегда подставлял дефолт TrainConfigGTAUAV.mona_bottleneck=64. Для StripNet поле mona_bottleneck (и mona_last_n_blocks, последнее не используется для StripNet) не вынесено в StripNetModelsConfig — это технический долг, отмечен ниже. Пока хардкод 64 совпадает с прежним поведением.

3.4 — Ветка `dinov3`

AsymmetricEncoder(
    dino_web_path=models_cfg.dino_web_path,
    dino_sat_path=models_cfg.dino_sat_path,
    lrsclip_path=models_common_cfg.lrsclip_path,
    init_gate=models_common_cfg.init_gate,
    baseline_mode=models_common_cfg.baseline_mode,
    shared_encoder=models_cfg.shared_encoder,
    mona_bottleneck=models_cfg.mona_bottleneck,
    mona_last_n_blocks=models_cfg.mona_last_n_blocks,
    device=hardware_cfg.device,
    backbone=backbone,
    stripnet_path="nn_models/STRIPNET/stripnet_s.pth",                         # старый дефолт TrainConfigGTAUAV
    stripnet_mona_last_n_stages=0,
    stripnet_freeze=True,
)

Почему stripnet_path передаётся для DINOv3: симметричная ситуация. AsymmetricEncoder.__init__ принимает параметр всегда, для DINOv3 он игнорируется. Старый код передавал cfg.stripnet_path (дефолт nn_models/STRIPNET/stripnet_s.pth) даже при DINOv3 — воспроизводим то же.

3.5 — Resume через `AsymmetricEncoder.load_checkpoint`

load_checkpoint принимает только 4 параметра (path, dino_web_path, dino_sat_path, lrsclip_path, device) — остальные восстанавливаются из чекпоинта. Передаём dino_*_path исходя из типа models_cfg:

DINOv3ModelsConfig → значения из конфига
StripNetModelsConfig → дефолтные значения _DINO_WEB/_DINO_SAT

Это уже было в старом коде; новый код просто аккуратнее распределил значения по типам конфига.

Известное ограничение (унаследовано от старого кода): AsymmetricEncoder.load_checkpoint не поддерживает StripNet-чекпоинты — он не принимает backbone='stripnet' и потому при resume StripNet-эксперимента построит DINOv3-модель. Это не регрессия — старый код имел тот же баг. Чинить — отдельный шаг.

4. Переименование `"sofia"` → `"sofia_v71"` в if-ах

В исходном коде разные части использовали разные имена для одного и того же бэкбона:

В if-ах сборки модели: if cfg.backbone == "sofia_v71"
В чекпоинт-блоке: if cfg.backbone in ("sofia", "sofia_v1") ← остаточное старое имя
В сообщении gradient_checkpointing: if cfg.backbone in ("stripnet", "sofia", "sofia_v1") ← тоже

Это был остаточный баг после промежуточного переименования sofia → sofia_v71 — в одних местах сделали, в других забыли. На уровне runtime это не приводило к падению (sofia-эксперименты тогда не запускались), но при первом запуске sofia-пресета чекпоинт-блок не сохранил бы sofia_cfg для v7.1 (ветка просто не сработала бы — backbone == "sofia_v71" не in ("sofia", "sofia_v1")).

Что сделано:

В новых presets/<name>/models.gin для Sofia v7.1: ModelsCommonConfig.backbone = 'sofia_v71'

В новом train_gtauav.py все if-ы используют "sofia_v71":

if backbone == "sofia_v71":              # сборка модели + resume + enc_str
if backbone in ("sofia_v71", "sofia_v1"):  # чекпоинт-блок (исправлено)
if backbone in ("stripnet", "sofia_v71", "sofia_v1"):  # gradient_checkpointing (исправлено)

В config_loader.py мапинг _BACKBONE_TO_MODELS_CLS: "sofia_v71": SOFIAv71ModelsConfig

Имена теперь согласованы на всех уровнях:

src/conf/models_common_conf.py → backbone: str ('dinov3' | 'stripnet' | 'sofia_v1' | 'sofia_v71')
src/conf/config_loader.py → словарь маппинга
presets/<name>/models.gin → биндинги
src/training/train_gtauav.py → все if-ы
src/models/sofia_v71/ → имя директории моделей

"sofia" без версии больше нигде не используется.

5. Файлы, которые добавились / изменились

Создано

src/main.py — точка входа
src/utils/__init__.py
src/utils/path_utils.py — get_proj_dir()
src/utils/seed_utils.py — set_seed() (на 4б)
src/utils/io_utils.py — atomic_save_torch(), clear_vram() (на 4б)

Перезаписано

src/training/train_gtauav.py — новый файл (см. train_gtauav.py в outputs)
src/losses/multi_infonce.py — снято @gin.configurable с InfoNCELoss (две строки удалены)
src/losses/weighted_infonce.py — снято @gin.configurable с WeightedInfoNCELoss (две строки удалены)

Удалено

conf/ (директория, 17 файлов) — старые .gin мёртвый код, биндят несуществующий TrainConfigGTAUAV

Не тронуто на 4а

src/datasets/gtauav_dataset.py — _RGB_ROOT/_CAPTION_ROOT остаются на module-level, но никто их теперь не использует. Удалить можно отдельным мини-коммитом или в 4б.
src/datasets/visloc_with_captions.py (legacy v2) — оставлен по решению пользователя.

6. Технический долг (на 4б)

mona_bottleneck для StripNet — вынести из хардкода 64 в StripNetModelsConfig.mona_bottleneck или в ModelsCommonConfig
Декомпозиция train() на Trainer + методы (1100 строк → ~50 строк за метод)
_evaluate → src/eval/evaluator.py с @torch.inference_mode() вместо @torch.no_grad()
CSVLogger → src/training/csv_logger.py
_atomic_save → atomic_save_torch из src/utils/io_utils.py (с cleanup .tmp на ошибке)
_set_seed / _clear_vram заменить на set_seed / clear_vram из src/utils/
AsymmetricEncoder.load_checkpoint для StripNet — расширить сигнатуру или сделать отдельный путь resume
Удалить _RGB_ROOT/_CAPTION_ROOT из gtauav_dataset.py

7. Контрольный smoke-test

После применения шага 4а:

cd <proj_dir>

# 1. Конфиги загружаются.
python -c "
from src.conf.config_loader import load_all_configs
from src.utils.path_utils import get_proj_dir
cfgs = load_all_configs(get_proj_dir() + 'in/config_files/', 'gtauav_balanced')
print('OK', sorted(cfgs.keys()), cfgs['models_common'].backbone)
"

# 2. Тренировка стартует.
python -m src.main gtauav_balanced

# 3. На 1 эпохе и 16 батчах метрики r@1_q2g/r@5_q2g/loss/tau/gate_q/gate_g
#    совпадают со старым запуском до 4-го знака.

Если все три проверки проходят — шаг 4а закрыт, можно идти в 4б.

16 KiB Raw Blame History Unescape Escape

Шаг 4а — Что изменилось

1. Новая точка входа — src/main.py

2. Изменённый src/training/train_gtauav.py