claude_refactor_v3: New train_gtauav.py, added entry point main.py, added utils
This commit is contained in:
291
belka_refactor_04_05_log.md
Normal file
291
belka_refactor_04_05_log.md
Normal file
@@ -0,0 +1,291 @@
|
|||||||
|
# Шаг 4а — Что изменилось
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 1. Новая точка входа — `src/main.py`
|
||||||
|
|
||||||
|
Запуск тренировки переехал из `src/training/train_gtauav.py::main()` в отдельный модуль `src/main.py`.
|
||||||
|
|
||||||
|
**Старый запуск:**
|
||||||
|
```bash
|
||||||
|
python src/training/train_gtauav.py --config conf/gtauav_balanced.gin
|
||||||
|
```
|
||||||
|
|
||||||
|
**Новый запуск:**
|
||||||
|
```bash
|
||||||
|
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()`
|
||||||
|
|
||||||
|
Было:
|
||||||
|
```python
|
||||||
|
def train(cfg: TrainConfigGTAUAV) -> None:
|
||||||
|
```
|
||||||
|
|
||||||
|
Стало:
|
||||||
|
```python
|
||||||
|
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:
|
||||||
|
|
||||||
|
```python
|
||||||
|
# было
|
||||||
|
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`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
# стало
|
||||||
|
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 убран
|
||||||
|
|
||||||
|
```python
|
||||||
|
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`
|
||||||
|
|
||||||
|
```python
|
||||||
|
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`
|
||||||
|
|
||||||
|
```python
|
||||||
|
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`
|
||||||
|
|
||||||
|
```python
|
||||||
|
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`
|
||||||
|
|
||||||
|
```python
|
||||||
|
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"`:
|
||||||
|
```python
|
||||||
|
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б)
|
||||||
|
|
||||||
|
1. **`mona_bottleneck` для StripNet** — вынести из хардкода `64` в `StripNetModelsConfig.mona_bottleneck` или в `ModelsCommonConfig`
|
||||||
|
2. **Декомпозиция `train()`** на `Trainer` + методы (1100 строк → ~50 строк за метод)
|
||||||
|
3. **`_evaluate` → `src/eval/evaluator.py`** с `@torch.inference_mode()` вместо `@torch.no_grad()`
|
||||||
|
4. **`CSVLogger` → `src/training/csv_logger.py`**
|
||||||
|
5. **`_atomic_save` → `atomic_save_torch`** из `src/utils/io_utils.py` (с cleanup `.tmp` на ошибке)
|
||||||
|
6. **`_set_seed` / `_clear_vram`** заменить на `set_seed` / `clear_vram` из `src/utils/`
|
||||||
|
7. **`AsymmetricEncoder.load_checkpoint`** для StripNet — расширить сигнатуру или сделать отдельный путь resume
|
||||||
|
8. **Удалить `_RGB_ROOT`/`_CAPTION_ROOT`** из `gtauav_dataset.py`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 7. Контрольный smoke-test
|
||||||
|
|
||||||
|
После применения шага 4а:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
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б.
|
||||||
@@ -10,7 +10,6 @@ Supports both learnable temperature (CLIP-style logit_scale) and fixed/scheduled
|
|||||||
|
|
||||||
import math
|
import math
|
||||||
|
|
||||||
import gin
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
import torch.nn.functional as F
|
import torch.nn.functional as F
|
||||||
@@ -81,10 +80,13 @@ def cosine_temperature(
|
|||||||
return tau_final + (tau_init - tau_final) * cosine
|
return tau_final + (tau_init - tau_final) * cosine
|
||||||
|
|
||||||
|
|
||||||
@gin.configurable
|
|
||||||
class InfoNCELoss(nn.Module):
|
class InfoNCELoss(nn.Module):
|
||||||
"""Symmetric InfoNCE with learnable or scheduled temperature.
|
"""Symmetric InfoNCE with learnable or scheduled temperature.
|
||||||
|
|
||||||
|
+ Note: NOT @gin.configurable. All parameters arrive explicitly from
|
||||||
|
+ train() via TrainingConfig — single source of truth for gin-bindable
|
||||||
|
+ values lives in src/conf/training_conf.py.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
temperature_init: Initial temperature value.
|
temperature_init: Initial temperature value.
|
||||||
temperature_final: Final temperature (only used if learnable=False).
|
temperature_final: Final temperature (only used if learnable=False).
|
||||||
|
|||||||
@@ -14,16 +14,17 @@ WeightedInfoNCE softens this with adaptive label smoothing per sample.
|
|||||||
|
|
||||||
import math
|
import math
|
||||||
|
|
||||||
import gin
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
import torch.nn.functional as F
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
|
||||||
@gin.configurable
|
|
||||||
class WeightedInfoNCELoss(nn.Module):
|
class WeightedInfoNCELoss(nn.Module):
|
||||||
"""Weighted InfoNCE with adaptive per-sample label smoothing.
|
"""Weighted InfoNCE with adaptive per-sample label smoothing.
|
||||||
|
|
||||||
|
+ Note: NOT @gin.configurable. All parameters arrive explicitly from
|
||||||
|
+ train() via TrainingConfig (loss_type='weighted' branch).
|
||||||
|
|
||||||
For each sample i, eps_i = 1 - (1 - base_smoothing) / (1 + exp(-k * w_i))
|
For each sample i, eps_i = 1 - (1 - base_smoothing) / (1 + exp(-k * w_i))
|
||||||
where w_i is the positive weight (e.g. IoU with matched satellite crop).
|
where w_i is the positive weight (e.g. IoU with matched satellite crop).
|
||||||
Higher weight → lower eps → sharper target (strong positive).
|
Higher weight → lower eps → sharper target (strong positive).
|
||||||
|
|||||||
54
src/main.py
Normal file
54
src/main.py
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
"""Entry point: load configs and run training.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python -m src.main gtauav_balanced
|
||||||
|
python -m src.main gtauav_balanced_sofia_v1
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import coloredlogs
|
||||||
|
|
||||||
|
from src.conf.config_loader import load_all_configs
|
||||||
|
from src.training.train_gtauav_old import train
|
||||||
|
from src.utils.path_utils import get_proj_dir
|
||||||
|
|
||||||
|
logger = logging.getLogger("caption_test")
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
coloredlogs.install(
|
||||||
|
level="INFO",
|
||||||
|
logger=logger,
|
||||||
|
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||||||
|
)
|
||||||
|
|
||||||
|
if len(sys.argv) != 2:
|
||||||
|
raise SystemExit(
|
||||||
|
"Usage: python -m src.main <preset_name>\n"
|
||||||
|
"Example: python -m src.main gtauav_balanced\n"
|
||||||
|
" available presets are subdirectories under in/config_files/",
|
||||||
|
)
|
||||||
|
preset_name = sys.argv[1]
|
||||||
|
|
||||||
|
proj_dir = get_proj_dir()
|
||||||
|
path2cfg = f"{proj_dir}in/config_files/" # per REQUIREMENTS_GIN_STYLE.md §5
|
||||||
|
|
||||||
|
configs = load_all_configs(path2cfg, preset_name)
|
||||||
|
|
||||||
|
train(
|
||||||
|
pipeline_cfg=configs["pipeline"],
|
||||||
|
hardware_cfg=configs["hardware"],
|
||||||
|
training_cfg=configs["training"],
|
||||||
|
tracking_cfg=configs["tracking"],
|
||||||
|
models_common_cfg=configs["models_common"],
|
||||||
|
models_cfg=configs["models"],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
|
||||||
File diff suppressed because it is too large
Load Diff
1395
src/training/train_gtauav_old.py
Normal file
1395
src/training/train_gtauav_old.py
Normal file
File diff suppressed because it is too large
Load Diff
2
src/utils/__init__.py
Normal file
2
src/utils/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
|||||||
|
"""Utilities: paths, seeding, IO."""
|
||||||
|
|
||||||
51
src/utils/io_utils.py
Normal file
51
src/utils/io_utils.py
Normal file
@@ -0,0 +1,51 @@
|
|||||||
|
"""IO helpers: atomic checkpoint saves, VRAM cleanup."""
|
||||||
|
|
||||||
|
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 (KeyboardInterrupt / SIGTERM included), the temp file
|
||||||
|
is removed. Makes --resume safe: a partial checkpoint never lands at
|
||||||
|
the destination path.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
obj: Anything torch.save can handle.
|
||||||
|
path: Destination path. Parent directory is created if missing.
|
||||||
|
"""
|
||||||
|
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."""
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
32
src/utils/path_utils.py
Normal file
32
src/utils/path_utils.py
Normal file
@@ -0,0 +1,32 @@
|
|||||||
|
"""Project root resolution via marker files."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Markers identifying the project root (per REQUIREMENTS_GIN_STYLE.md §5).
|
||||||
|
_MARKERS: tuple[str, ...] = ("pyproject.toml", ".git", "in")
|
||||||
|
|
||||||
|
|
||||||
|
def get_proj_dir() -> str:
|
||||||
|
"""Return absolute project root with trailing slash.
|
||||||
|
|
||||||
|
Walks up from this file's directory until finding pyproject.toml,
|
||||||
|
.git, or in/. Searches up to 10 levels.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Project root path 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}",
|
||||||
|
)
|
||||||
|
|
||||||
23
src/utils/seed_utils.py
Normal file
23
src/utils/seed_utils.py
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
"""RNG seeding for reproducibility."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import random
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
|
||||||
|
def set_seed(seed: int = 42) -> None:
|
||||||
|
"""Fix all RNG seeds (Python random, NumPy, PyTorch CPU + all CUDA devices).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
seed: Integer seed.
|
||||||
|
"""
|
||||||
|
random.seed(seed)
|
||||||
|
np.random.seed(seed)
|
||||||
|
torch.manual_seed(seed)
|
||||||
|
torch.cuda.manual_seed_all(seed)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
Reference in New Issue
Block a user