Compare commits
33 Commits
main
...
belka_refa
| Author | SHA1 | Date | |
|---|---|---|---|
|
ecf8071220 | ||
|
|
d82fac9190 | ||
|
|
7b3acc633a | ||
|
|
455ae2e99f | ||
|
|
1c03811af0 | ||
|
|
676dea7932 | ||
|
|
c9083c89cc | ||
|
|
9878473419 | ||
|
|
2362ce0adb | ||
|
|
911d2ce4e6 | ||
|
|
562a5e2e43 | ||
|
|
4cb5ab97d8 | ||
|
|
248bd331d2 | ||
|
|
4e148a29bb | ||
|
|
c43b4c82b9 | ||
|
|
80d1806cee | ||
|
|
2bef5a4270 | ||
|
|
7bbf67c465 | ||
|
|
89cb8ab0f7 | ||
|
|
08d328db09 | ||
|
|
4b441279e0 | ||
|
|
e8a0de7ad3 | ||
|
|
db2b5b32f4 | ||
|
|
af207d302c | ||
|
|
ce9e3dc94b | ||
|
|
bf5e417094 | ||
|
|
d42ef94821 | ||
|
|
6b928634f4 | ||
|
|
3275a41b36 | ||
|
|
4bbe1dda1d | ||
|
|
5c48b6c8fd | ||
|
|
0d8c82acc3 | ||
|
|
a27b5a7357 |
14
CLAUDE.md
14
CLAUDE.md
@@ -175,20 +175,20 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization.
|
||||
|
||||
## Датасет: GTA-UAV-LR
|
||||
|
||||
- **RGB:** `/home/servml/Документы/datasets/GTA-UAV-LR/`
|
||||
- **RGB:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR/`
|
||||
- Drone: 33,763 PNG (512x384), altitudes 100-600m
|
||||
- Satellite: 14,640 PNG (256x256 RGBA)
|
||||
- Pairs: `cross-area-drone2sate-{train,test}.json` (primary split)
|
||||
- Metadata: `*_drone_meta.csv` (height, yaw, roll, pitch)
|
||||
- Origin: GTA V simulation (Los Santos)
|
||||
|
||||
- **Captions:** `/home/servml/Документы/datasets/GTA-UAV-LR-captions/`
|
||||
- **Captions:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions/`
|
||||
- Drone: 33,411 JSON (32,635 multi-paragraph P1/P2/P3 + 776 short water-only)
|
||||
- Satellite: 6,546 JSON (все multi-paragraph)
|
||||
- Формат: 3 абзаца (P1 Inventory + P2 Spatial + P3 Fingerprint)
|
||||
- Token counts: ~430 output tokens per caption
|
||||
|
||||
- **Segmentation:** `/home/servml/Документы/datasets/GTA-UAV-LR-aug/`
|
||||
- **Segmentation:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/`
|
||||
- 48,403 images, 17 классов (background, building, road, vegetation, water, ...)
|
||||
- Modalities: segm/, depth/, edge/, chm/, safetensors/
|
||||
- Query: 512x512, DB: 256x256
|
||||
@@ -331,7 +331,7 @@ python -m scripts.compare_runs \
|
||||
|
||||
### UAV-VisLoc Prepare
|
||||
- **Путь:** `/home/servml/Документы/code/Yaroslav/UAV-VisLoc-prepare/scripts/prepare_dataset.py`
|
||||
- **Статус:** выполнен (2026-04-17), данные в `/home/servml/Документы/datasets/UAV_VisLoc_processed/` (25 GB)
|
||||
- **Статус:** выполнен (2026-04-17), данные в `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_processed/` (25 GB)
|
||||
- **Задача:** нарезка satellite кропов 512x512, stride 256 + resize drone -> 512x512
|
||||
- **Подробности:** см. ниже
|
||||
|
||||
@@ -376,7 +376,7 @@ Binary masks — natural FiLM gates. Modality dropout: text 0.3, CHM 0.5, seg 0.
|
||||
## Датасеты (справочник)
|
||||
|
||||
### UAV-VisLoc
|
||||
- **Путь:** `/home/servml/Документы/datasets/UAV_VisLoc_dataset/`
|
||||
- **Путь:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_dataset/`
|
||||
- **Структура:** 11 маршрутов (папки `01`-`11`), каждая содержит:
|
||||
- `drone/` — drone-снимки (`XX_NNNN.JPG`)
|
||||
- `satelliteXX.tif` — спутниковая карта
|
||||
@@ -403,8 +403,8 @@ Binary masks — natural FiLM gates. Modality dropout: text 0.3, CHM 0.5, seg 0.
|
||||
### Запуск
|
||||
```bash
|
||||
python scripts/prepare_dataset.py \
|
||||
--src /home/servml/Документы/datasets/UAV_VisLoc_dataset \
|
||||
--dst /home/servml/Документы/datasets/UAV_VisLoc_processed \
|
||||
--src /media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_dataset \
|
||||
--dst /media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_processed \
|
||||
--crop-size 512 --stride 256 --target-size 512
|
||||
```
|
||||
|
||||
|
||||
158
DIAGNOSTIC_DSS_MoCo_2026-04-25.md
Normal file
158
DIAGNOSTIC_DSS_MoCo_2026-04-25.md
Normal file
@@ -0,0 +1,158 @@
|
||||
# Диагностика: коллапс recall на эпохе 1 — DSS + MoCo queue
|
||||
|
||||
**Дата:** 2026-04-25
|
||||
**Конфиг:** `conf/gtauav_balanced.gin` (gate=0.7, with text)
|
||||
**Запуск:** `python -m src.training.train_gtauav --config conf/gtauav_balanced.gin --filter-meta meta/seg_filter.json`
|
||||
|
||||
---
|
||||
|
||||
## 1. Симптомы из лога
|
||||
|
||||
### Метрики по эпохам
|
||||
|
||||
| Эпоха | Режим сэмплера | LR | train loss | train R@1 | val R@1 | val AP | eval loss |
|
||||
|---|---|---|---|---|---|---|---|
|
||||
| **ep0** | mutex-only (warmup) | 5.00e-05 | 4.2239 | **0.0742** | 0.0758 | 0.1511 | 2.0783 |
|
||||
| **ep1** | DSS активирован, MoCo queue полная | 1.00e-04 | 4.4218 ↑ | **0.0120** ↓6× | 0.0107 | 0.0328 | 2.0786 |
|
||||
|
||||
### Ключевые наблюдения
|
||||
- **Recall обвалился ровно на эпохе включения DSS** (warmup=1, поэтому DSS стартует на ep1)
|
||||
- **Train loss выросла** (4.22 → 4.42), хотя должна падать
|
||||
- **Eval loss практически стоит** (2.0783 → 2.0786) — модель не учится на чистом батч-сигнале
|
||||
- **Разрыв train/eval loss > 2×** — большая часть train loss идёт от MoCo+hard-mining компонент, не от in-batch contrast
|
||||
- Gate stable: `gate_q=0.70`, `gate_g=0.70` (текст не вытаскивает дополнительный градиент)
|
||||
- `tau=0.07` (clamped, learnable, но почти не двигается → loss не учит температуру)
|
||||
|
||||
### Сбой при сохранении чекпойнта
|
||||
```
|
||||
RuntimeError: [enforce fail at inline_container.cc:668] .
|
||||
unexpected pos 289048128 vs 289048016
|
||||
```
|
||||
**Причина:** диск переполнен (ENOSPC). Подтверждено независимо: `mkdir` в session-env тоже падает с ENOSPC. Это инфраструктурный баг, не алгоритмический — но он не даёт сохранить состояние и продолжить.
|
||||
|
||||
---
|
||||
|
||||
## 2. Почему DSS «не работает» на этом этапе
|
||||
|
||||
### 2.1. Mutex-constraint спасает от false negatives только формально
|
||||
Mutex исключает дронов с пересекающимися `sat_candidates`. Но DSS специально пакует **визуально похожих** дронов — у которых sat-кандидаты могут быть *разными tiles*, но при этом сцены практически идентичные:
|
||||
- один и тот же район Лос-Сантоса
|
||||
- та же высота полёта
|
||||
- та же time-of-day, та же погода (синтетика GTA-V)
|
||||
|
||||
Для холодного энкодера такие «negatives» неотличимы от позитива → softmax пытается развести то, что развести нечем → градиенты шумные и противоречивые.
|
||||
|
||||
### 2.2. Re-embed раз в эпоху — слишком редко в начале обучения
|
||||
- На ep1 LR удваивается (5e-5 → 1e-4 после warmup)
|
||||
- Encoder за эпоху смещается сильно (MONA-адаптеры учатся быстро)
|
||||
- Батчи внутри эпохи собираются по embeddings, посчитанным **до** этих обновлений
|
||||
- То есть «similarity» в DSS — это similarity **вчерашнего** энкодера
|
||||
- Re-embed на ep1 занял 313s, но между двумя re-embed encoder успевает измениться существенно
|
||||
|
||||
### 2.3. Эпохи 0 на mutex-only недостаточно как warmup
|
||||
- К моменту запуска DSS R@1=0.074 на полной gallery (~2684 уникальных тайлов)
|
||||
- Random baseline ≈ 1/2684 ≈ 0.0004
|
||||
- 0.074 — это «модель чуть-чуть оторвалась от случая», но качество embeddings ещё недостаточно
|
||||
- DSS усиливает шум, а не сигнал: «визуально похожих» определяет *шумный* энкодер
|
||||
|
||||
---
|
||||
|
||||
## 3. Почему MoCo queue делает хуже
|
||||
|
||||
### 3.1. Representation drift в queue
|
||||
- Queue хранит **4096 embeddings**, посчитанных в разные моменты времени разными версиями encoder
|
||||
- На ep0–1, когда MONA-адаптеры учатся быстро, разница между «свежим» и «3 шага назад» embedding'ом — заметная
|
||||
- Эти устаревшие негативы дают сигнал в направлении, в котором энкодер уже не находится
|
||||
- Без momentum encoder (как в оригинальном MoCo) drift ничем не сглажен
|
||||
|
||||
### 3.2. `hard_mining_k=512` амплифицирует ошибку queue
|
||||
- Из 4096 берутся 512 «самых трудных»
|
||||
- «Трудность» меряется в **текущем** feature space
|
||||
- Сами вектора лежали там в **старом** feature space
|
||||
- На холодном энкодере это эквивалентно «учиться отличать себя-позитива от себя-вчерашнего»
|
||||
- Loss растёт, recall падает
|
||||
|
||||
### 3.3. Тройная hard-negative композиция → collapse-режим
|
||||
Все три механизма независимо подсовывают «трудные» негативы:
|
||||
| Механизм | Источник трудности |
|
||||
|---|---|
|
||||
| DSS | визуально близкие в батче |
|
||||
| Mutex sampler | in-batch contrast после фильтра |
|
||||
| `hard_mining_k=512` | top-K из MoCo queue |
|
||||
|
||||
На warm encoder это polish; на cold encoder это **too-hard negatives problem** (известная failure mode у contrastive learning, Robinson et al., 2021):
|
||||
- Модель не получает «лёгких» примеров, на которых формируется базовое embedding-пространство
|
||||
- Градиенты толкают её в произвольных направлениях
|
||||
- Embedding space коллапсирует или размывается
|
||||
|
||||
---
|
||||
|
||||
## 4. Подтверждения из истории коммитов
|
||||
|
||||
```
|
||||
8f8cbb1 Diagnostic baseline v2: also disable MoCo queue
|
||||
c25bd64 Diagnostic baseline: disable DSS + hard mining, fresh output dir
|
||||
9a7fbff Fix plot_combined: fallback from 'total' to 'train_loss'
|
||||
70f1617 Fix autograd in-place error: move memory-bank enqueue after backward
|
||||
8197ab2 Fix training loop: only pass positive_weights to WeightedInfoNCELoss
|
||||
```
|
||||
|
||||
Подозрение на DSS+MoCo уже было — текущие числа подтверждают его эмпирически.
|
||||
|
||||
---
|
||||
|
||||
## 5. Рекомендации
|
||||
|
||||
### 5.1. Cold-start curriculum (приоритет)
|
||||
1. **Несколько эпох mutex-only без MoCo и без hard_mining** до R@1 хотя бы ~0.2 на train.
|
||||
- Конкретно: `sampler_type="mutex"`, `loss.use_memory_bank=False`, `loss.hard_mining_k=0`.
|
||||
2. Затем **включить MoCo queue с warm-up**: либо momentum encoder, либо queue pre-fill 1 эпоху со `stop_grad` на queue updates.
|
||||
3. **DSS включать только после** того, как embeddings стали discriminative (R@1 на train > 0.2). Иначе «similar» = «random».
|
||||
4. **`hard_mining_k` стартует с 0** и поднимается curriculum-схемой (например, 0 → 64 → 256 → 512 по эпохам).
|
||||
|
||||
### 5.2. Изменения по DSS
|
||||
- Сократить интервал re-embed (раз в N шагов, не раз в эпоху) — минимум первые 2-3 эпохи
|
||||
- Или временно фиксировать кеш embeddings из чистого baseline (без DSS) и использовать его как «референсный» для сэмплинга
|
||||
|
||||
### 5.3. Изменения по MoCo
|
||||
- Добавить momentum encoder (EMA на ключи, как в оригинальном MoCo) — это решает drift
|
||||
- Либо очищать queue на каждой эпохе (теряем эффект, но избегаем drift)
|
||||
- Размер queue 4096 при batch 64 → 64 батча в очереди = слишком долгая история для холодного энкодера
|
||||
|
||||
### 5.4. Инфраструктура
|
||||
- **Освободить диск** (ENOSPC блокирует чекпойнты и часть утилит)
|
||||
- Возможно, перенести `out/gtauav/` на другой диск или почистить старые runs
|
||||
- Добавить pre-flight disk check перед `torch.save`
|
||||
|
||||
---
|
||||
|
||||
## 6. План эксперимента для подтверждения
|
||||
|
||||
| Run | Sampler | MoCo queue | hard_mining_k | Цель |
|
||||
|---|---|---|---|---|
|
||||
| **A** (clean baseline) | mutex | off | 0 | Подтвердить, что без DSS/MoCo recall растёт нормально |
|
||||
| **B** (only MoCo) | mutex | on | 0 | Изолировать вклад queue |
|
||||
| **C** (only hard_mining) | mutex | on | 512 | Изолировать вклад top-K mining |
|
||||
| **D** (only DSS) | dss | off | 0 | Изолировать вклад DSS |
|
||||
| **E** (full, current) | dss | on | 512 | Reproduce coллапс |
|
||||
|
||||
Decision rule:
|
||||
- Если A учится нормально (R@1 растёт монотонно), а E коллапсирует — это decisive ablation
|
||||
- Если B/C/D по отдельности тоже коллапсируют — проблема в каждом компоненте
|
||||
- Если только E коллапсирует — проблема в композиции
|
||||
|
||||
---
|
||||
|
||||
## 7. Альтернативная гипотеза (менее вероятная)
|
||||
|
||||
LR=1e-4 на проекциях после warmup может быть слишком большим для shared DINOv3 + MONA с малым trainable %. Симптомы похожи (резкий обвал на эпохе с полным LR), но не объясняют, почему mutex-only baseline (commit `c25bd64`) учится нормально без изменения LR. Так что это вторичный фактор.
|
||||
|
||||
---
|
||||
|
||||
## 8. Резюме
|
||||
|
||||
**Главный диагноз:** связка `DSS + MoCo queue + hard_mining_k=512` создаёт **too-hard negatives problem** на холодном энкодере. Каждый компонент по отдельности рассчитан на warm encoder; их композиция на ep1 (когда R@1 ещё ~0.07) делает задачу нерешаемой и приводит к коллапсу embedding space.
|
||||
|
||||
**Решение:** curriculum — сначала mutex+CE без queue/mining до R@1≈0.2, потом постепенно включать остальные механизмы.
|
||||
|
||||
**Блокер:** диск переполнен (ENOSPC), без освобождения места дальнейшее обучение и сохранение чекпойнтов невозможно.
|
||||
431
GTA-UAV-LR_Dataset_Analysis.md
Normal file
431
GTA-UAV-LR_Dataset_Analysis.md
Normal file
@@ -0,0 +1,431 @@
|
||||
# АНАЛИЗ ДАТАСЕТА: GTA-UAV (LR)
|
||||
|
||||
**Дата анализа:** 2026-04-21
|
||||
**Метод:** Эмпирический анализ данных на диске + статья arXiv:2409.16925 + GitHub-репозиторий авторов
|
||||
**Путь к данным:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR/`
|
||||
**Путь к аугментациям:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/`
|
||||
|
||||
---
|
||||
|
||||
## 1. МЕТАДАННЫЕ
|
||||
|
||||
| Поле | Значение |
|
||||
|------|----------|
|
||||
| Полное название | Game4Loc: A UAV Geo-Localization Benchmark from Game Data |
|
||||
| Авторы | Yuxiang Ji, Boyong He, Zhuoyue Tan, Liaoni Wu |
|
||||
| Год, Venue | 2025, AAAI 2025 (Oral), arXiv:2409.16925 [cs.CV] |
|
||||
| Код | https://github.com/Yux1angJi/GTA-UAV (Python, PyTorch) |
|
||||
| Данные | HuggingFace / BaiduDisk (пароль: gtav) |
|
||||
| Лицензия | Apache 2.0 |
|
||||
| Общий объём на диске | **~26 GB** (LR); HR-версия ~143 GB |
|
||||
|
||||
---
|
||||
|
||||
## 2. ОБЩАЯ СТАТИСТИКА
|
||||
|
||||
### 2.1. Сводка
|
||||
|
||||
| Параметр | Значение |
|
||||
|----------|----------|
|
||||
| Drone-изображений (query) | **33 763** |
|
||||
| Спутниковых тайлов (DB) | **14 640** (4 уровня зума) |
|
||||
| Всего изображений | **48 403** |
|
||||
| Покрытая территория | **81.3 км²** (один непрерывный регион) |
|
||||
| Сцена | Лос-Сантос (GTA V) — городская + природная среда |
|
||||
| Номинальных высот | 6 (100, 200, 300, 400, 500, 600 м) |
|
||||
| Направлений камеры (yaw) | **361 уникальное** значение (произвольные углы) |
|
||||
| Тип данных | **Синтетика** (рендеринг в игровом движке GTA V) |
|
||||
|
||||
### 2.2. Разбиения
|
||||
|
||||
Датасет предоставляет **два протокола** оценки:
|
||||
|
||||
| Протокол | Train | Test | Метод разбиения |
|
||||
|----------|-------|------|-----------------|
|
||||
| Same-area | 26 964 (80.0%) | 6 744 (20.0%) | По изображениям в пределах одного региона |
|
||||
| Cross-area | 15 693 (46.5%) | 18 015 (53.5%) | По географическим подрегионам |
|
||||
|
||||
**Same-area:** train и test из одной и той же территории, без пересечений по изображениям.
|
||||
**Cross-area:** train и test из непересекающихся географических зон — более сложный и реалистичный сценарий.
|
||||
|
||||
### 2.3. Распределение по номинальным высотам
|
||||
|
||||
| Номинальная высота | Same-area train | Same-area test | Всего | Доля |
|
||||
|-------------------|----------------|---------------|-------|------|
|
||||
| 100 м | 4 248 | 1 072 | 5 320 | 15.8% |
|
||||
| 200 м | 4 572 | 1 162 | 5 734 | 17.0% |
|
||||
| 300 м | 4 546 | 1 138 | 5 684 | 16.9% |
|
||||
| 400 м | 4 615 | 1 162 | 5 777 | 17.1% |
|
||||
| 500 м | 4 598 | 1 171 | 5 769 | 17.1% |
|
||||
| 600 м | 4 385 | 1 039 | 5 424 | 16.1% |
|
||||
| **Итого** | **26 964** | **6 744** | **33 708** | 100% |
|
||||
|
||||
Распределение по высотам **равномерное** (~16–17% на каждую номинальную высоту).
|
||||
|
||||
---
|
||||
|
||||
## 3. ИСТОЧНИКИ ИЗОБРАЖЕНИЙ
|
||||
|
||||
### 3.1. Дроновые виды (query)
|
||||
|
||||
| Параметр | Значение |
|
||||
|----------|----------|
|
||||
| Платформа | **Синтетический рендеринг** (GTA V / DeepGTAV) |
|
||||
| Тип съёмки | RGB, near-nadir (камера ≈90° вниз, cam_roll ≈ -90°) |
|
||||
| Разрешение (LR) | **512x384 px** |
|
||||
| Разрешение (HR) | 1920x1440 px |
|
||||
| Pitch камеры | ≈-90° (вертикально вниз), drone_pitch ±10° |
|
||||
| Roll дрона | ±20° (реалистичные вибрации) |
|
||||
| Yaw (heading) | Произвольные углы 0°–360° (361 уникальное значение) |
|
||||
| Формат | PNG |
|
||||
|
||||
### 3.2. Спутниковые виды (DB)
|
||||
|
||||
| Параметр | Значение |
|
||||
|----------|----------|
|
||||
| Платформа | Спутниковые тайлы (стиль Google Maps) |
|
||||
| Формат | PNG |
|
||||
| Разрешение тайлов | **256x256 px** (все уровни зума) |
|
||||
| Количество тайлов | **14 640** |
|
||||
| Уровни зума | 4 уровня (4, 5, 6, 7) |
|
||||
|
||||
### 3.3. Мультимасштабная тайловая сетка
|
||||
|
||||
Спутниковые данные организованы как **пирамида зум-уровней** — каждый следующий уровень удваивает разрешение:
|
||||
|
||||
| Зум | Cols x Rows | Тайлов | Масштаб (относительно zoom 7) |
|
||||
|-----|-------------|--------|------------------------------|
|
||||
| 4 | 11 x 16 | 176 | 1/8 |
|
||||
| 5 | 22 x 32 | 704 | 1/4 |
|
||||
| 6 | 43 x 64 | 2 752 | 1/2 |
|
||||
| 7 | 86 x 128 | 11 008 | 1/1 (базовый) |
|
||||
| **Итого** | | **14 640** | |
|
||||
|
||||
Все сетки заполнены на **100%** (нет пропущенных тайлов).
|
||||
|
||||
**Общий размер карты** (zoom 7): 86 × 256 = **22 016 px** по ширине, 128 × 256 = **32 768 px** по высоте.
|
||||
|
||||
Именование тайлов: `{zoom}_{0}_{X}_{Y}.png` (пример: `7_0_42_63.png`).
|
||||
|
||||
---
|
||||
|
||||
## 4. ПАРАМЕТРЫ СЪЁМКИ ДРОНОВ
|
||||
|
||||
### 4.1. Фактические высоты полёта
|
||||
|
||||
Номинальные высоты (из имени файла) **не совпадают** точно с фактическими — это реалистичный эффект рельефа:
|
||||
|
||||
| Номинальная | Фактический диапазон | Среднее | Кол-во |
|
||||
|-------------|---------------------|---------|--------|
|
||||
| 100 м | 24.9 – 444.9 м | 93.9 м | 5 320 |
|
||||
| 200 м | 34.3 – 548.4 м | 193.9 м | 5 734 |
|
||||
| 300 м | 31.0 – 648.9 м | 294.3 м | 5 684 |
|
||||
| 400 м | 33.4 – 746.0 м | 395.2 м | 5 777 |
|
||||
| 500 м | 44.9 – 930.4 м | 494.9 м | 5 769 |
|
||||
| 600 м | 82.1 – 988.1 м | 593.6 м | 5 424 |
|
||||
|
||||
**Общий диапазон фактических высот:** 24.9 – 988.1 м (среднее 345.3 м).
|
||||
|
||||
Большой разброс объясняется тем, что высота задаётся **над уровнем моря**, а рельеф варьируется.
|
||||
|
||||
### 4.2. Ориентация камеры (attitude)
|
||||
|
||||
| Параметр | Среднее | Диапазон |
|
||||
|----------|---------|----------|
|
||||
| Drone roll (|abs|) | 3.99° | 0° – 20° |
|
||||
| Drone pitch (|abs|) | 3.87° | 0° – 10° |
|
||||
| Cam roll | -90.0° | -110° – -71.2° |
|
||||
| Drone yaw | произвольный | -180° – +180° (361 уникальное значение) |
|
||||
|
||||
Камера направлена **вертикально вниз** (cam_roll ≈ -90°), с реалистичными отклонениями дрона по roll/pitch.
|
||||
|
||||
### 4.3. Классификация углов (для обучения)
|
||||
|
||||
В CSV-файлах `*_class_angle.csv` yaw дискретизирован в **4 класса** (квадранты):
|
||||
|
||||
| Класс | Диапазон | Кол-во (cross-area train) |
|
||||
|-------|----------|--------------------------|
|
||||
| 0 | 0°–90° | 7 871 |
|
||||
| 1 | 90°–180° | 7 637 |
|
||||
| 2 | 180°–270° | 7 650 |
|
||||
| 3 | 270°–360° | 7 834 |
|
||||
|
||||
Распределение по квадрантам **равномерное**.
|
||||
|
||||
---
|
||||
|
||||
## 5. ПРОСТРАНСТВЕННОЕ ПОКРЫТИЕ
|
||||
|
||||
### 5.1. Зона дронов в координатах карты
|
||||
|
||||
| Параметр | Значение |
|
||||
|----------|----------|
|
||||
| Диапазон X (zoom 7) | 606 – 6 068 px (охват ~5 462 px) |
|
||||
| Диапазон Y (zoom 7) | 1 330 – 9 735 px (охват ~8 405 px) |
|
||||
| Покрытие карты | 5 462 / 22 016 × 8 405 / 32 768 ≈ **6.4%** площади zoom-7 карты |
|
||||
|
||||
Дрон-изображения сконцентрированы в определённом подрегионе карты, при этом **вся** карта (100% тайлов) доступна как DB.
|
||||
|
||||
### 5.2. Регион
|
||||
|
||||
| Параметр | Значение |
|
||||
|----------|----------|
|
||||
| Источник | Карта Лос-Сантос (GTA V) |
|
||||
| Площадь | 81.3 км² |
|
||||
| Типы ландшафта | Город, пригороды, горы, побережье, леса, пустыня |
|
||||
| Координатный охват | Синтетические координаты (x, y в пикселях карты) |
|
||||
|
||||
---
|
||||
|
||||
## 6. АННОТАЦИИ И МЕТАДАННЫЕ
|
||||
|
||||
### 6.1. Файлы аннотаций
|
||||
|
||||
| Файл | Содержание | Формат |
|
||||
|------|-----------|--------|
|
||||
| `same-area-drone2sate-train.json` | Same-area train: 26 964 записей | JSON (list of dicts) |
|
||||
| `same-area-drone2sate-test.json` | Same-area test: 6 744 записей | JSON |
|
||||
| `cross-area-drone2sate-train.json` | Cross-area train: 15 693 записей | JSON |
|
||||
| `cross-area-drone2sate-test.json` | Cross-area test: 18 015 записей | JSON |
|
||||
| `*_drone_meta.csv` | Метаданные дрона (height, yaw, roll, pitch) | TSV |
|
||||
| `*_drone_meta_new.csv` | Обновлённые метаданные | TSV |
|
||||
| `*_class_angle.csv` | Класс угла yaw (0–3) | CSV |
|
||||
| `back_csv/` | Обратные CSV-файлы для cross-area | Директория |
|
||||
|
||||
### 6.2. Структура JSON-записи
|
||||
|
||||
```json
|
||||
{
|
||||
"drone_img_dir": "drone/images",
|
||||
"drone_img_name": "400_0001_0000022427.png",
|
||||
"drone_loc_x_y": [1268.6, 6326.8],
|
||||
"sate_img_dir": "satellite",
|
||||
"pair_pos_sate_img_list": ["5_0_3_18.png"],
|
||||
"pair_pos_sate_weight_list": [0.521],
|
||||
"pair_pos_sate_loc_x_y_list": [[1209.6, 6393.6]],
|
||||
"pair_pos_semipos_sate_img_list": ["4_0_1_9.png", "5_0_3_17.png", ...],
|
||||
"pair_pos_semipos_sate_weight_list": [0.174, 0.146, ...],
|
||||
"pair_pos_semipos_sate_loc_x_y_list": [[1036.8, 6566.4], ...],
|
||||
"drone_metadata": {
|
||||
"height": 356.92,
|
||||
"drone_roll": -3.66,
|
||||
"drone_pitch": -4.81,
|
||||
"drone_yaw": -86.0,
|
||||
"cam_roll": -93.66,
|
||||
"cam_pitch": -4.81,
|
||||
"cam_yaw": -86.0
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 6.3. Типы аннотаций
|
||||
|
||||
| Тип аннотации | Наличие | Комментарий |
|
||||
|---------------|---------|-------------|
|
||||
| Координаты drone (x, y) | **Да** | В пространстве карты (zoom 7) |
|
||||
| Высота дрона (altitude) | **Да** | В метрах, 6 номинальных уровней |
|
||||
| Heading angle (yaw) | **Да** | Произвольные углы -180°...+180° |
|
||||
| Pitch / Roll | **Да** | Drone и Camera отдельно |
|
||||
| Positive pairs | **Да** | С весами IoU overlap |
|
||||
| Semi-positive pairs | **Да** | Множественные, с весами |
|
||||
| Координаты тайлов (x, y) | **Да** | Центры тайлов в пространстве карты |
|
||||
| GPS-координаты | **Нет** | Синтетическая карта |
|
||||
| Depth maps | Нет | — (но сгенерированы в aug) |
|
||||
| Segmentation masks | Нет | — (но сгенерированы в aug) |
|
||||
| Bounding boxes | Нет | — |
|
||||
| Временная метка | Нет | — |
|
||||
|
||||
---
|
||||
|
||||
## 7. РАСПРЕДЕЛЕНИЕ POSITIVE / SEMI-POSITIVE ПАР
|
||||
|
||||
### 7.1. Positive matches (точные)
|
||||
|
||||
| Positives на query | Кол-во | Доля |
|
||||
|-------------------|--------|------|
|
||||
| 0 | 16 297 | 48.3% |
|
||||
| 1 | 16 505 | 49.0% |
|
||||
| 2 | 906 | 2.7% |
|
||||
|
||||
**~48% query не имеют точного positive** — это ключевая особенность датасета. Авторы моделируют реалистичный сценарий, когда дрон-изображение может не иметь точно совпадающего спутникового тайла.
|
||||
|
||||
### 7.2. Semi-positive matches (частичное перекрытие)
|
||||
|
||||
| Semi-positives | Кол-во | Доля |
|
||||
|----------------|--------|------|
|
||||
| 1 | 93 | 0.3% |
|
||||
| 2 | 1 538 | 4.6% |
|
||||
| 3 | 7 301 | 21.7% |
|
||||
| 4 | 7 290 | 21.6% |
|
||||
| 5 | 8 135 | 24.1% |
|
||||
| 6 | 6 685 | 19.8% |
|
||||
| 7 | 2 168 | 6.4% |
|
||||
| 8 | 369 | 1.1% |
|
||||
| 9–10 | 129 | 0.4% |
|
||||
|
||||
**Среднее число semi-positives на query:** 4.58
|
||||
|
||||
В отличие от UAV-GeoLoc (где ~94% query имеют ровно 4 positive), здесь распределение гораздо более **вариативное** (от 1 до 10 semi-positives).
|
||||
|
||||
### 7.3. Уникальные спутниковые тайлы в парах
|
||||
|
||||
Из 14 640 спутниковых тайлов, в positive/semi-positive парах задействовано **4 266** уникальных тайлов (29.1%).
|
||||
|
||||
Распределение по зумам среди используемых:
|
||||
|
||||
| Зум | Тайлов в парах | Из общего числа |
|
||||
|-----|---------------|-----------------|
|
||||
| 4 | 83 | 47.2% |
|
||||
| 5 | 279 | 39.6% |
|
||||
| 6 | 886 | 32.2% |
|
||||
| 7 | 3 018 | 27.4% |
|
||||
|
||||
---
|
||||
|
||||
## 8. АУГМЕНТИРОВАННЫЙ НАБОР (GTA-UAV-LR-aug)
|
||||
|
||||
**Путь:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/`
|
||||
**Объём:** ~71 GB
|
||||
|
||||
### 8.1. Сгенерированные модальности
|
||||
|
||||
| Модальность | Модель | Drone | Satellite | Итого |
|
||||
|-------------|--------|-------|-----------|-------|
|
||||
| Depth | DA3-LARGE-1.1 | 33 763 | 14 640 | 48 403 |
|
||||
| CHM (Canopy Height) | DINOv3-ViTL16-CHMv2 | 33 763 | 14 640 | 48 403 |
|
||||
| Edge | Sobel from depth (CPU) | 33 763 | 14 640 | 48 403 |
|
||||
| Segmentation | SegEarth-OV3 | 33 763 | 14 640 | 48 403 |
|
||||
|
||||
### 8.2. Параметры пайплайна
|
||||
|
||||
| Параметр | Значение |
|
||||
|----------|----------|
|
||||
| Версия пайплайна | 4.0.0-dir-layout |
|
||||
| Размер DB-изображений | 256 px |
|
||||
| Размер query-изображений | 512 px |
|
||||
| Профиль GPU | RTX 4090 |
|
||||
| Сохранение safetensors | Да |
|
||||
| Дата генерации | 2026-04-19 |
|
||||
|
||||
### 8.3. Классы сегментации (17 промптов)
|
||||
|
||||
background, building, road, vegetation, water, sand and gravel ground, rocky terrain, farmland, railway, parking lot, sidewalk, bare soil and plowed field, roof and rooftop, sports field and playground, muddy ground and wetland, embankment and levee, swimming pool
|
||||
|
||||
### 8.4. Дополнительные форматы
|
||||
|
||||
| Директория | Содержание |
|
||||
|-----------|-----------|
|
||||
| `npy/` | Предвычисленные NumPy-массивы (chm, depth, edge, segm) |
|
||||
| `safetensors/` | Эмбеддинги (drone, satellite) |
|
||||
|
||||
---
|
||||
|
||||
## 9. СРАВНЕНИЕ С UAV-GeoLoc И UAV-VisLoc
|
||||
|
||||
| Параметр | GTA-UAV (LR) | UAV-GeoLoc | UAV-VisLoc |
|
||||
|----------|-------------|-----------|-----------|
|
||||
| Тип данных | **Синтетика** (GTA V) | **Синтетика** (Google Earth Studio) | **Реальные** БПЛА |
|
||||
| Drone-изображений | 33 763 | 652 744 | 6 774 |
|
||||
| DB-изображений | 14 640 | 274 683 | 0 (нужно нарезать) |
|
||||
| Всего изображений | 48 403 | 927 427 | 6 774 + 11 карт |
|
||||
| Разрешение drone | **512x384** | 512x512 | 3976x2652 / 3000x2000 |
|
||||
| Разрешение DB | **256x256** | 100x100 – 1000x1000 | N/A (целые карты) |
|
||||
| Высоты полёта | 25–988 м (6 номинальных) | 100, 125, 150 м | 405–2572 м |
|
||||
| Heading (yaw) | Произвольные (361 значение) | 8 × 45° | Произвольные |
|
||||
| Регионы | 1 (Лос-Сантос, GTA V) | 372 сцены, 11 стран | 11 маршрутов, 7 провинций Китая |
|
||||
| Сцен/маршрутов | 1 непрерывная карта | 372 | 11 |
|
||||
| Площадь покрытия | 81.3 км² | Варьируется | Варьируется |
|
||||
| Positive pairs | Да (с весами IoU) | Да (positive.json) | Нет (нужно по GPS) |
|
||||
| Semi-positive pairs | Да (с весами IoU) | Да (semi_positive.json) | Нет |
|
||||
| % query без positive | **48.3%** | ~1.2% | — |
|
||||
| Avg semi-positives | 4.58 | ~2.83 | — |
|
||||
| Split протоколы | Same-area + Cross-area | Terrain / Country / All | По изображениям (~75/25) |
|
||||
| Мультимасштабность DB | **Да** (4 зум-уровня) | Нет (фиксированный размер) | Нет |
|
||||
| Temporal gap | **Нет** (одновременно) | Нет | 2–5 лет |
|
||||
| Лицензия | Apache 2.0 | CC BY-NC 4.0 | Не указана |
|
||||
| Объём на диске | ~26 GB (LR) | ~181 GB | ~16.4 GB |
|
||||
|
||||
### Ключевые отличия GTA-UAV:
|
||||
|
||||
1. **Непрерывная территория** — одна большая карта 81.3 км², а не набор дискретных сцен. Это позволяет моделировать реальный сценарий навигации.
|
||||
2. **Мультимасштабные тайлы** — 4 уровня зума (4–7) vs единый размер кропов в UAV-GeoLoc.
|
||||
3. **Частичные совпадения** — 48% query не имеют точного positive, что моделирует реалистичный сценарий, когда дрон находится между тайлами. Используются **веса IoU** вместо бинарных меток.
|
||||
4. **Широкий диапазон высот** — 6 номинальных уровней (100–600 м) vs 3 уровня (100–150 м) в UAV-GeoLoc.
|
||||
5. **Произвольные углы** — 361 уникальное значение yaw vs 8 дискретных (шаг 45°) в UAV-GeoLoc.
|
||||
6. **Два протокола оценки** — same-area (проще) и cross-area (реалистичнее) vs terrain/country/all splits в UAV-GeoLoc.
|
||||
|
||||
---
|
||||
|
||||
## 10. ПОЛНАЯ СТРУКТУРА ДАННЫХ
|
||||
|
||||
```
|
||||
GTA-UAV-LR/ # ~26 GB
|
||||
├── README.md # Описание датасета
|
||||
│
|
||||
├── drone/
|
||||
│ └── images/
|
||||
│ ├── 100_0001_0000000000.png # 512x384
|
||||
│ ├── 100_0001_0000000001.png
|
||||
│ ├── ...
|
||||
│ ├── 200_0001_XXXXXXXXXX.png
|
||||
│ ├── ...
|
||||
│ └── 600_0001_XXXXXXXXXX.png # 33 763 файлов
|
||||
│
|
||||
├── satellite/
|
||||
│ ├── 4_0_0_0.png # 256x256, zoom 4
|
||||
│ ├── 4_0_0_1.png
|
||||
│ ├── ... # 176 тайлов zoom 4
|
||||
│ ├── 5_0_0_0.png # 704 тайлов zoom 5
|
||||
│ ├── ...
|
||||
│ ├── 6_0_0_0.png # 2 752 тайлов zoom 6
|
||||
│ ├── ...
|
||||
│ ├── 7_0_0_0.png # 11 008 тайлов zoom 7
|
||||
│ └── ... # 14 640 файлов итого
|
||||
│
|
||||
├── same-area-drone2sate-train.json # 26 964 записей
|
||||
├── same-area-drone2sate-test.json # 6 744 записей
|
||||
├── cross-area-drone2sate-train.json # 15 693 записей
|
||||
├── cross-area-drone2sate-test.json # 18 015 записей
|
||||
│
|
||||
├── cross-area-drone2sate-train_drone_meta.csv # TSV: img, height, yaw, roll, pitch
|
||||
├── cross-area-drone2sate-test_drone_meta.csv
|
||||
├── cross-area-drone2sate-train_drone_meta_new.csv # Обновлённые метаданные
|
||||
├── cross-area-drone2sate-test_drone_meta_new.csv
|
||||
├── cross-area-drone2sate-train_class_angle.csv # CSV: img, angle_class (0-3)
|
||||
├── cross-area-drone2sate-test_class_angle.csv
|
||||
│
|
||||
└── back_csv/ # Обратные CSV для cross-area
|
||||
├── cross-area-drone2sate-train_class_angle.csv
|
||||
└── cross-area-drone2sate-test_class_angle.csv
|
||||
|
||||
|
||||
GTA-UAV-LR-aug/ # ~71 GB (аугментации)
|
||||
├── manifest.json # Метаданные пайплайна
|
||||
├── pipeline.log
|
||||
│
|
||||
├── depth/ # DA3-LARGE-1.1
|
||||
│ ├── drone/images/ (33 763 файлов)
|
||||
│ └── satellite/ (14 640 файлов)
|
||||
│
|
||||
├── chm/ # DINOv3-ViTL16-CHMv2
|
||||
│ ├── drone/images/ (33 763 файлов)
|
||||
│ └── satellite/ (14 640 файлов)
|
||||
│
|
||||
├── edge/ # Sobel from depth
|
||||
│ ├── drone/images/ (33 763 файлов)
|
||||
│ └── satellite/ (14 640 файлов)
|
||||
│
|
||||
├── segm/ # SegEarth-OV3, 17 классов
|
||||
│ ├── drone/images/ (33 763 файлов)
|
||||
│ └── satellite/ (14 640 файлов)
|
||||
│
|
||||
├── npy/ # NumPy массивы
|
||||
│ ├── chm/
|
||||
│ ├── depth/
|
||||
│ ├── edge/
|
||||
│ └── segm/
|
||||
│
|
||||
└── safetensors/ # Предвычисленные эмбеддинги
|
||||
├── drone/
|
||||
└── satellite/
|
||||
```
|
||||
24
README.md
24
README.md
@@ -377,9 +377,9 @@ that list it as a valid candidate. `n_scored_g2q` is reported in metrics for tra
|
||||
### V3 — GTA-UAV + DINOv3 + DGTRS-CLIP (active)
|
||||
|
||||
**Dataset:** GTA-UAV-LR (33K drone + 14K satellite, GTA V synthetic)
|
||||
- RGB: `/home/servml/Документы/datasets/GTA-UAV-LR/`
|
||||
- Captions: `/home/servml/Документы/datasets/GTA-UAV-LR-captions/` (40K JSON, 3-paragraph VLM)
|
||||
- Segmentation: `/home/servml/Документы/datasets/GTA-UAV-LR-aug/` (17 classes)
|
||||
- RGB: `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR/`
|
||||
- Captions: `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions/` (40K JSON, 3-paragraph VLM)
|
||||
- Segmentation: `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/` (17 classes)
|
||||
- Seg filter: 37,498 passed / 10,905 excluded (>=90% background+water)
|
||||
- Split: 80/20 random (26,966 train / 6,742 test → 24,891/6,252 after seg filter)
|
||||
|
||||
@@ -584,3 +584,21 @@ tensorboard --logdir out/gtauav/with_text/tb_logs
|
||||
- Type hints on all signatures
|
||||
- Google-style docstrings
|
||||
- English-only comments
|
||||
|
||||
# Authors
|
||||
|
||||
## Лицензия и автор
|
||||
|
||||
**Автор:** Пикалёв Ярослав Сергеевич, к.т.н., ЛИСАД, ФГБНУ «ИПИИ»
|
||||
**Email:** i@pikaliov.ru
|
||||
|
||||
**Рефакторинг ветка:** Павленко Богдан Викторович, м.н.с., ЛИСАД, ФГБНУ «ИПИИ»
|
||||
**Email:** bogdanpavl2000@mail.ru
|
||||
|
||||
Проект реализован в рамках: (НИР FREN-2024-0002 «Извлечение семантической информации из изображений для автономных систем навигации БПЛА».)
|
||||
|
||||
==**Код предоставляется для научных целей в рамках выполняемой работы.**==
|
||||
|
||||
**Работы, ссылка на которые обязательна при использовании этой реализации:**
|
||||
1. Пикалёв Я. С., Павленко Б. В. Регрессионная нейронная сеть на основе регуляризации функции потерь… // Информатика и автоматизация (SPCRAS). 2025.
|
||||
2. Павленко Б. В., Пикалёв Я. С. Методика создания набора аэрофотоснимков для CVGL // ПИИ. 2024. № 4 (35). С. 101–112.
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
**Дата анализа:** 2026-04-17
|
||||
**Метод:** Эмпирический анализ данных на диске + статья arXiv:2405.11936 + GitHub-репозиторий авторов
|
||||
**Путь к данным:** `/home/servml/Документы/datasets/UAV_VisLoc_dataset/`
|
||||
**Путь к данным:** `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/UAV_VisLoc_dataset/`
|
||||
|
||||
---
|
||||
|
||||
|
||||
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б.
|
||||
208
belka_refactor_07_05_log.md
Normal file
208
belka_refactor_07_05_log.md
Normal file
@@ -0,0 +1,208 @@
|
||||
# Шаг 4b — Что изменилось
|
||||
|
||||
---
|
||||
|
||||
## 1. Новая точка входа — `src/main.py` + добавлена debug-конфигурация `Main entry point` в `launch.json`
|
||||
|
||||
**Пример запуска отладки:**
|
||||
```bash
|
||||
python src/training/train_gtauav.py --config conf/gtauav_balanced.gin
|
||||
```
|
||||
|
||||
## 2. Тесты
|
||||
|
||||
Добавлены `tests\conftest.py & test_trainer.py`
|
||||
|
||||
### Тест загрузки конфигураций в трейнер
|
||||
Файл `test_trainer.py`
|
||||
|
||||
**Таблица конфигураций**
|
||||
|
||||
- scr/conf/*.py конфиги
|
||||
- in/config_files общие common .gin-конфиги и пресеты под эксперименты
|
||||
---
|
||||
|
||||
| Пресет | backbone | baseline_mode | shared_encoder | init_gate |
|
||||
| ----------------------------------- | --------- | ------------- | -------------- | --------- |
|
||||
| `gtauav_balanced` | dinov3 | False | True | 0.7 |
|
||||
| `gtauav_baseline` | dinov3 | True | True | 0.7 |
|
||||
| `gtauav_balanced_asym` | dinov3 | False | False | 0.7 |
|
||||
| `gtauav_baseline_asym` | dinov3 | True | False | 0.7 |
|
||||
| `gtauav_text_heavy` | dinov3 | False | True | **0.3** |
|
||||
| `gtauav_image_heavy` | dinov3 | False | True | **0.9** |
|
||||
| `gtauav_balanced_stripnet` | stripnet | False | — | 0.7 |
|
||||
| `gtauav_balanced_stripnet_unfrozen` | stripnet | False | — | 0.7 |
|
||||
| `gtauav_baseline_stripnet` | stripnet | True | — | 0.7 |
|
||||
| `gtauav_baseline_stripnet_unfrozen` | stripnet | True | — | 0.7 |
|
||||
| `gtauav_balanced_sofia` | sofia_v71 | False | — | 0.7 |
|
||||
| `gtauav_baseline_sofia` | sofia_v71 | True | — | 0.7 |
|
||||
| `gtauav_balanced_sofia_v1` | sofia_v1 | False | — | 0.7 |
|
||||
| `gtauav_baseline_sofia_v1` | sofia_v1 | True | — | 0.7 |
|
||||
| `preprocess` | — | — | — | — |
|
||||
|
||||
|
||||
**Пример запуска отладки:**
|
||||
```bash
|
||||
python -m pytest tests/test_trainer.py
|
||||
```
|
||||
|
||||
**Лог:**
|
||||
|
||||
```bash
|
||||
=========================================== test session starts ===========================================
|
||||
platform linux -- Python 3.10.19, pytest-9.0.3, pluggy-1.6.0
|
||||
rootdir: /home/servml/Документы/code/Yaroslav/caption-test
|
||||
collected 21 items
|
||||
|
||||
tests/test_trainer.py . ✅ _SUPPORTED_BACKBONES must be a frozenset (immutable, hashable).
|
||||
. ✅ Exactly dinov3 and stripnet are supported in the current refactor.
|
||||
. ✅ ModelsConfig union mirrors _SUPPORTED_BACKBONES (dinov3 | stripnet).
|
||||
.
|
||||
Supported backbones pass _validate_backbone silently.
|
||||
✅ dinov3
|
||||
. ✅ stripnet
|
||||
.
|
||||
Unsupported backbones (incl. sofia) raise NotImplementedError, not ImportError.
|
||||
✅ sofia_v1
|
||||
. ✅ sofia_v71
|
||||
. ✅ mistral_42b
|
||||
. ✅ empty
|
||||
.
|
||||
Trainer(...) instantiates from every real preset's loaded cfgs.
|
||||
✅ gtauav_balanced
|
||||
. ✅ gtauav_balanced_asym
|
||||
. ✅ gtauav_baseline
|
||||
. ✅ gtauav_baseline_asym
|
||||
. ✅ gtauav_image_heavy
|
||||
. ✅ gtauav_text_heavy
|
||||
. ✅ gtauav_balanced_stripnet
|
||||
. ✅ gtauav_balanced_stripnet_unfrozen
|
||||
. ✅ gtauav_baseline_stripnet
|
||||
. ✅ gtauav_baseline_stripnet_unfrozen
|
||||
. ✅ All runtime fields are None / 0 / [] before .train() is called.
|
||||
. ✅ Trainer.train() takes only `self` — main.py calls trainer.train().
|
||||
|
||||
[100%]
|
||||
|
||||
=========================================== 21 passed in 2.19s ===========================================
|
||||
```
|
||||
|
||||
## 3. Декомпозиция вызовов в `trainer_new.py`
|
||||
|
||||
**Теперь вызовы выглядят такой цепочкой:**
|
||||
```python
|
||||
def train(self) -> None:
|
||||
"""Full pipeline: setup → build → train → evaluate → cleanup."""
|
||||
self._validate_backbone()
|
||||
clear_vram()
|
||||
set_seed(self.pipeline_cfg.seed)
|
||||
self._setup_output_dir()
|
||||
self._setup_tracker()
|
||||
self._build_model()
|
||||
self._configure_gradient_checkpointing()
|
||||
self._log_model_summary()
|
||||
self._build_loss()
|
||||
self._build_neg_bank()
|
||||
self._build_data_loaders()
|
||||
self._build_optimizer_and_scheduler()
|
||||
self._restore_from_resume()
|
||||
self._setup_profiler()
|
||||
|
||||
try:
|
||||
self._train_loop()
|
||||
self._final_evaluation()
|
||||
finally:
|
||||
self._cleanup()
|
||||
```
|
||||
|
||||
**Лог:**
|
||||
|
||||
### Тест вызовов до блока `try: self._train_loop()`
|
||||
✅ Загрузка весов
|
||||
✅ Сборка пайплайна
|
||||
✅ Загрузка моделей
|
||||
✅ Загрузка данных
|
||||
|
||||
```bash
|
||||
2026-05-07 09:40:53 caption_test.trainer INFO ⚙️ Validate backbone
|
||||
2026-05-07 09:40:53 caption_test.trainer INFO ⚙️ Setup out dir
|
||||
2026-05-07 09:40:53 caption_test.trainer INFO ⚙️ Setup tracker...
|
||||
2026-05-07 09:40:53 caption_test.trackers WARNING tensorboard not installed, skipping TB tracking
|
||||
2026-05-07 09:41:24 caption_test.trainer INFO ⚙️ Build loss...
|
||||
2026-05-07 09:41:52 caption_test.trainer INFO Building model — with text (L1/L2/L3), shared DINOv3 WEB
|
||||
2026-05-07 09:42:24 caption_test.trainer INFO ⚙️ Build DINOv3 model...
|
||||
2026-05-07 09:50:48 caption_test.model INFO 🧊 Loading DINOv3 from dinov3-vitl16-pretrain-lvd1689m.pth
|
||||
2026-05-07 09:50:48 caption_test.model INFO 🧊 Loading DINOv3 from dinov3-vitl16-pretrain-lvd1689m.pth
|
||||
2026-05-07 09:54:48 caption_test.model INFO 🧊 DINOv3 loaded: 303,129,600 params
|
||||
2026-05-07 09:54:48 caption_test.model INFO 🧊 DINOv3 loaded: 303,129,600 params
|
||||
2026-05-07 09:55:03 caption_test.adapters INFO 🔧 MONA injected: 24 adapters (blocks 12-23 of 24), 3,502,080 trainable params (bottleneck=64)
|
||||
2026-05-07 09:55:03 caption_test.adapters INFO 🔧 MONA injected: 24 adapters (blocks 12-23 of 24), 3,502,080 trainable params (bottleneck=64)
|
||||
2026-05-07 09:55:05 caption_test.model INFO Shared encoder mode: single DINOv3 WEB for drone + satellite
|
||||
2026-05-07 09:55:05 caption_test.model INFO Shared encoder mode: single DINOv3 WEB for drone + satellite
|
||||
2026-05-07 09:55:38 caption_test.dgtrs INFO 📝 Loading DGTRS-CLIP text encoder from DGTRS-CLIP-ViT-L-14.pt
|
||||
2026-05-07 09:55:38 caption_test.dgtrs INFO 📝 Loading DGTRS-CLIP text encoder from DGTRS-CLIP-ViT-L-14.pt
|
||||
2026-05-07 09:55:39 caption_test.dgtrs INFO 📝 DGTRS text encoder loaded: 123,972,096 params, context=248 tokens
|
||||
2026-05-07 09:55:39 caption_test.dgtrs INFO 📝 DGTRS text encoder loaded: 123,972,096 params, context=248 tokens
|
||||
2026-05-07 10:03:01 caption_test.adapters INFO 🔧 LoRA injected: 12 blocks, rank=4, 147,456 trainable params
|
||||
2026-05-07 10:03:01 caption_test.adapters INFO 🔧 LoRA injected: 12 blocks, rank=4, 147,456 trainable params
|
||||
2026-05-07 10:04:01 caption_test.trainer INFO embed_dim=1024
|
||||
2026-05-07 10:04:22 caption_test.trainer INFO ⚙️ Configure gradient checkpointing...
|
||||
2026-05-07 10:04:42 caption_test.trainer INFO ✅ Gradient checkpointing enabled (DINOv3 + DGTRS)
|
||||
2026-05-07 10:04:51 caption_test.trainer INFO trainable=7,059,458 (1.63%) total=434,161,154
|
||||
2026-05-07 10:04:51 caption_test.profiler INFO torchinfo not installed, using basic parameter count
|
||||
2026-05-07 10:04:51 caption_test.profiler INFO Model summary:
|
||||
Total parameters: 434,161,154
|
||||
Trainable parameters: 7,059,458 (1.63%)
|
||||
[trainable] image_encoder.layer.12.mona_attn.gamma: [1024] (1,024)
|
||||
... (Вывод списка параметров)
|
||||
...
|
||||
2026-05-07 10:05:00 caption_test.trainer INFO ⚙️ Build loss...
|
||||
2026-05-07 10:06:06 caption_test.trainer INFO Loss: SymmetricInfoNCE Temperature: learnable (init=0.070) q2g=0.60 g2q=0.40
|
||||
2026-05-07 10:06:07 caption_test.trainer INFO ⚙️ Build negative bank...
|
||||
2026-05-07 10:06:13 caption_test.trainer INFO ⚙️ Build dataloaders...
|
||||
2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
|
||||
2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
|
||||
2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
|
||||
2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
|
||||
2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
|
||||
2026-05-07 10:06:13 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
|
||||
2026-05-07 10:06:19 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
|
||||
2026-05-07 10:06:19 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
|
||||
2026-05-07 10:06:20 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO ✅ Loaded 6252 pairs from meta/test_20.json
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO ✅ Loaded 6252 pairs from meta/test_20.json
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 🔻 Filter loaded: 10905 excluded images
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📐 Altitude index: 33708 drones (from 2 CSV)
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
|
||||
2026-05-07 10:06:22 caption_test.gtauav_dataset INFO 📚 Loading caption index from /home/servml/Документы/datasets/GTA-UAV-LR-captions
|
||||
2026-05-07 10:06:24 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
|
||||
2026-05-07 10:06:24 caption_test.gtauav_dataset INFO 📚 Caption index: 39957 entries
|
||||
2026-05-07 10:06:25 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
|
||||
2026-05-07 10:06:25 caption_test.gtauav_dataset INFO ✅ Loaded 24891 pairs from meta/train_80.json
|
||||
2026-05-07 10:06:25 caption_test.trainer INFO Sampler: MutuallyExclusive — no false negatives within a batch
|
||||
2026-05-07 10:06:25 caption_test.trainer INFO train=24891 test=6252 batch=8 accum=8 effective_batch=64
|
||||
2026-05-07 10:06:46 caption_test.trainer INFO ⚙️ Build optimizer & scheduler...
|
||||
2026-05-07 10:06:46 caption_test.trainer INFO Optimizer: AdamW LR: proj=1e-04 text=1e-05 warmup=2 epochs
|
||||
2026-05-07 10:07:33 caption_test.trainer INFO ⚙️ Restore from resume...
|
||||
2026-05-07 10:07:37 caption_test.trainer INFO ⚙️ Setup profiler...
|
||||
```
|
||||
|
||||
`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(...)` именованными аргументами
|
||||
|
||||
|
||||
@@ -1,48 +0,0 @@
|
||||
# Balanced: GatedFusion with text captions enabled.
|
||||
# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs satellite
|
||||
|
||||
import src.datasets.visloc_with_captions
|
||||
import src.losses.multi_infonce
|
||||
import src.models.dual_encoder
|
||||
import src.training.train
|
||||
|
||||
# ---- Dual encoder ----
|
||||
DualEncoderCaptionTest.variant = "ViT-B-32"
|
||||
DualEncoderCaptionTest.pretrained_path = "checkpoints/RS5M_ViT-B-32.pt"
|
||||
DualEncoderCaptionTest.unfreeze_mode = "last_block"
|
||||
DualEncoderCaptionTest.embed_dim = 512
|
||||
DualEncoderCaptionTest.use_mlp_heads = False
|
||||
DualEncoderCaptionTest.baseline_mode = False
|
||||
DualEncoderCaptionTest.init_gate = 0.7
|
||||
DualEncoderCaptionTest.device = "cuda"
|
||||
|
||||
# ---- Fusion ----
|
||||
GatedFusion.init_gate = 0.7
|
||||
GatedFusion.baseline_mode = False
|
||||
|
||||
# ---- Loss ----
|
||||
InfoNCELoss.temperature_init = 0.1
|
||||
InfoNCELoss.temperature_final = 0.01
|
||||
InfoNCELoss.label_smoothing = 0.1
|
||||
InfoNCELoss.weight_q2g = 0.6
|
||||
InfoNCELoss.weight_g2q = 0.4
|
||||
|
||||
# ---- Dataset ----
|
||||
GeoLocCaptionDataset.drop_caption_prob = 0.0
|
||||
GeoLocCaptionDataset.seed = 42
|
||||
|
||||
# ---- Training ----
|
||||
TrainConfig.train_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/train_query.txt"
|
||||
TrainConfig.val_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/val_query.txt"
|
||||
TrainConfig.data_root = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc"
|
||||
TrainConfig.output_dir = "out/caption_test/balanced"
|
||||
TrainConfig.epochs = 10
|
||||
TrainConfig.batch_size = 128
|
||||
TrainConfig.num_workers = 4
|
||||
TrainConfig.learning_rate = 1e-4
|
||||
TrainConfig.weight_decay = 1e-4
|
||||
TrainConfig.grad_clip = 1.0
|
||||
TrainConfig.use_amp = True
|
||||
TrainConfig.eval_every = 2
|
||||
TrainConfig.seed = 42
|
||||
TrainConfig.device = "cuda"
|
||||
@@ -1,10 +0,0 @@
|
||||
# Baseline: no text fusion (gate forced to 1.0).
|
||||
# query = drone_only -> InfoNCE vs satellite
|
||||
# Reference R@1 for delta computation.
|
||||
|
||||
include 'balanced.gin'
|
||||
|
||||
DualEncoderCaptionTest.baseline_mode = True
|
||||
GatedFusion.baseline_mode = True
|
||||
|
||||
TrainConfig.output_dir = "out/caption_test/baseline_no_text"
|
||||
@@ -1,64 +0,0 @@
|
||||
# GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions.
|
||||
# Symmetric InfoNCE + MutuallyExclusiveSampler (no false negatives).
|
||||
# 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank.
|
||||
|
||||
import src.losses.multi_infonce
|
||||
|
||||
# ---- Training ----
|
||||
TrainConfigGTAUAV.epochs = 10
|
||||
TrainConfigGTAUAV.batch_size = 8
|
||||
TrainConfigGTAUAV.num_workers = 4
|
||||
TrainConfigGTAUAV.learning_rate = 1e-4
|
||||
TrainConfigGTAUAV.text_lr_factor = 0.1
|
||||
TrainConfigGTAUAV.weight_decay = 1e-4
|
||||
TrainConfigGTAUAV.grad_clip = 1.0
|
||||
TrainConfigGTAUAV.grad_accum_steps = 8
|
||||
TrainConfigGTAUAV.use_amp = True
|
||||
TrainConfigGTAUAV.eval_every = 1
|
||||
TrainConfigGTAUAV.warmup_epochs = 2
|
||||
TrainConfigGTAUAV.seed = 42
|
||||
TrainConfigGTAUAV.device = "cuda"
|
||||
|
||||
# ---- Model ----
|
||||
TrainConfigGTAUAV.init_gate = 0.7
|
||||
TrainConfigGTAUAV.baseline_mode = False
|
||||
TrainConfigGTAUAV.shared_encoder = True # single DINOv3 WEB for both branches
|
||||
TrainConfigGTAUAV.mona_bottleneck = 64
|
||||
TrainConfigGTAUAV.mona_last_n_blocks = 12 # inject MONA only in last 12/24 ViT blocks
|
||||
TrainConfigGTAUAV.gradient_checkpointing = True
|
||||
|
||||
# ---- Loss ----
|
||||
TrainConfigGTAUAV.loss_type = "symmetric"
|
||||
TrainConfigGTAUAV.tau_init = 0.07
|
||||
TrainConfigGTAUAV.label_smoothing = 0.1
|
||||
TrainConfigGTAUAV.learnable_temperature = True
|
||||
TrainConfigGTAUAV.weight_q2g = 0.6
|
||||
TrainConfigGTAUAV.weight_g2q = 0.4
|
||||
TrainConfigGTAUAV.neg_bank_size = 4096
|
||||
|
||||
# ---- Sampling ----
|
||||
TrainConfigGTAUAV.sampler_type = "dss" # "dss" or "mutex"
|
||||
TrainConfigGTAUAV.dss_warmup_epochs = 1 # first N epochs use mutex-only (untrained embeds not useful)
|
||||
TrainConfigGTAUAV.dss_reembed_every = 1
|
||||
TrainConfigGTAUAV.use_mutex_sampler = True # legacy flag, kept True unless disabling both samplers
|
||||
|
||||
# ---- Output ----
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"
|
||||
|
||||
# ---- Tracking ----
|
||||
TrainConfigGTAUAV.use_wandb = False
|
||||
TrainConfigGTAUAV.use_tb = True
|
||||
TrainConfigGTAUAV.use_gradcam = True
|
||||
TrainConfigGTAUAV.gradcam_every = 5
|
||||
TrainConfigGTAUAV.use_profiler = False
|
||||
TrainConfigGTAUAV.log_grad_norms = True
|
||||
|
||||
# ---- InfoNCELoss (gin-configurable) ----
|
||||
InfoNCELoss.temperature_init = 0.07
|
||||
InfoNCELoss.learnable_temperature = True
|
||||
InfoNCELoss.label_smoothing = 0.1
|
||||
InfoNCELoss.weight_q2g = 0.6
|
||||
InfoNCELoss.weight_g2q = 0.4
|
||||
InfoNCELoss.tau_min = 0.01
|
||||
InfoNCELoss.tau_max = 0.1
|
||||
InfoNCELoss.hard_mining_k = 512 # 0 = use whole queue (disable mining)
|
||||
@@ -1,16 +0,0 @@
|
||||
# GTA-UAV Balanced (asymmetric, full MONA): WEB drone encoder + SAT satellite encoder.
|
||||
# MONA injected into all 24 ViT blocks of each encoder.
|
||||
# Same loss/sampling/optimizer as gtauav_balanced.gin; differs only in model arch.
|
||||
#
|
||||
# Trainable: ~17.6M (MONA 2× × 24 blocks + LoRA + TextFusionMLP + gates + tau)
|
||||
# Total params: ~733M (2× DINOv3-L + DGTRS-CLIP)
|
||||
# VRAM target (RTX 4090, 24 GB): ~16-20 GB at batch=8 with gradient checkpointing.
|
||||
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
# ---- Model overrides: asymmetric + full MONA ----
|
||||
TrainConfigGTAUAV.shared_encoder = False # WEB for drone, SAT for satellite
|
||||
TrainConfigGTAUAV.mona_last_n_blocks = 24 # MONA in all 24 ViT blocks (was 12)
|
||||
|
||||
# ---- Output ----
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/balanced_asym"
|
||||
@@ -1,27 +0,0 @@
|
||||
# GTA-UAV Balanced (StripNet backbone): StripNet-small + Conv-MONA in last 2 stages.
|
||||
# Replaces DINOv3 ViT-L/16 with strip-shaped DWConv hierarchical CNN (~28M params,
|
||||
# 10× smaller than DINOv3). Output 512-dim → projected to 1024 to match retrieval space.
|
||||
#
|
||||
# Trainable:
|
||||
# - Projection (Linear 512→1024): ~525K
|
||||
# - Conv-MONA in stages 3 & 4 (2 adapters per Block × 6 blocks total): ~2-3M
|
||||
# - LoRA on DGTRS-CLIP: 147K
|
||||
# - TextFusionMLP (shared): ~3.4M
|
||||
# - GatedFusion gates + tau: 3 scalars
|
||||
#
|
||||
# StripNet pretrained on ImageNet-1K (head dropped); state-dict naming follows
|
||||
# upstream Strip-R-CNN repo (`conv_spatial1/2`).
|
||||
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
# ---- Backbone ----
|
||||
TrainConfigGTAUAV.backbone = "stripnet"
|
||||
TrainConfigGTAUAV.stripnet_path = "nn_models/STRIPNET/stripnet_s.pth"
|
||||
TrainConfigGTAUAV.stripnet_mona_last_n_stages = 2 # Conv-MONA in stages 3 & 4 (deepest)
|
||||
|
||||
# ---- Model overrides ----
|
||||
TrainConfigGTAUAV.shared_encoder = True # StripNet always shared (one CNN for both branches)
|
||||
TrainConfigGTAUAV.mona_bottleneck = 64 # Conv-MONA bottleneck channels
|
||||
|
||||
# ---- Output ----
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_stripnet"
|
||||
@@ -1,17 +0,0 @@
|
||||
# GTA-UAV Balanced (StripNet, fully unfrozen): all StripNet layers trainable.
|
||||
# Backbone trains with reduced LR (lr * stripnet_backbone_lr_factor).
|
||||
# Conv-MONA disabled by default — full fine-tune supplies enough capacity.
|
||||
# Set stripnet_mona_last_n_stages > 0 if you want MONA + fine-tune hybrid.
|
||||
#
|
||||
# Note: StripNet uses BatchNorm. With small batch (8) and gradient accumulation,
|
||||
# BN running stats may drift. Watch validation loss for instability.
|
||||
|
||||
include 'conf/gtauav_balanced_stripnet.gin'
|
||||
|
||||
# ---- Unfreeze backbone ----
|
||||
TrainConfigGTAUAV.stripnet_freeze = False
|
||||
TrainConfigGTAUAV.stripnet_mona_last_n_stages = 0 # disable Conv-MONA (full fine-tune handles adaptation)
|
||||
TrainConfigGTAUAV.stripnet_backbone_lr_factor = 0.1 # backbone lr = 1e-4 * 0.1 = 1e-5
|
||||
|
||||
# ---- Output ----
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_stripnet_unfrozen"
|
||||
@@ -1,28 +0,0 @@
|
||||
# GTA-UAV Baseline: no text fusion (gate forced to 1.0).
|
||||
# query = drone_only -> InfoNCE vs satellite
|
||||
# Reference R@1 for delta computation.
|
||||
#
|
||||
# Diagnostic mode (2026-04-24): DSS and hard-negative mining disabled after
|
||||
# the previous run collapsed (R@1 = 0.6% at epoch 8, train loss growing).
|
||||
# Hypothesis: DSS packs visually-identical drones at bs=8 and the hard-mining
|
||||
# queue amplifies that hardness — together they prevent convergence from a
|
||||
# nearly-random start. Run with mutex-only sampling and the full queue as
|
||||
# uniform negatives first, restore the extras incrementally once baseline
|
||||
# converges.
|
||||
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
TrainConfigGTAUAV.baseline_mode = True
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_inbatch"
|
||||
TrainConfigGTAUAV.use_gradcam = False
|
||||
|
||||
# ---- Diagnostic overrides ----
|
||||
# Previous mutex-only run still collapsed at epoch 1 (val loss locked at log(8)).
|
||||
# Hypothesis refined: the MoCo-style queue stays stale because we have no
|
||||
# momentum encoder, and with reduced trainable surface (MONA-12) the model
|
||||
# can't reconcile fresh representations against 4096 stale negatives —
|
||||
# mode collapse. Disable the queue entirely so InfoNCE sees only the 8
|
||||
# fresh in-batch negatives, matching the OLD run's effective setup.
|
||||
TrainConfigGTAUAV.sampler_type = "mutex" # was "dss"
|
||||
TrainConfigGTAUAV.neg_bank_size = 0 # was 4096 — disable MoCo queue (no momentum encoder)
|
||||
InfoNCELoss.hard_mining_k = 0 # was 512 — irrelevant when queue is empty
|
||||
@@ -1,9 +0,0 @@
|
||||
# GTA-UAV Baseline (asymmetric, full MONA): no text fusion (gate forced to 1.0).
|
||||
# WEB drone encoder + SAT satellite encoder, MONA in all 24 ViT blocks.
|
||||
# Reference R@1 for delta computation against gtauav_balanced_asym.gin.
|
||||
|
||||
include 'conf/gtauav_balanced_asym.gin'
|
||||
|
||||
TrainConfigGTAUAV.baseline_mode = True
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_asym"
|
||||
TrainConfigGTAUAV.use_gradcam = False
|
||||
@@ -1,8 +0,0 @@
|
||||
# GTA-UAV Baseline (StripNet backbone): no text fusion. Reference R@1 for
|
||||
# computing Δ R@1 against gtauav_balanced_stripnet.gin.
|
||||
|
||||
include 'conf/gtauav_balanced_stripnet.gin'
|
||||
|
||||
TrainConfigGTAUAV.baseline_mode = True
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_stripnet"
|
||||
TrainConfigGTAUAV.use_gradcam = False
|
||||
@@ -1,8 +0,0 @@
|
||||
# GTA-UAV Baseline (StripNet, fully unfrozen): no text fusion. For Δ R@1
|
||||
# against gtauav_balanced_stripnet_unfrozen.gin.
|
||||
|
||||
include 'conf/gtauav_balanced_stripnet_unfrozen.gin'
|
||||
|
||||
TrainConfigGTAUAV.baseline_mode = True
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_stripnet_unfrozen"
|
||||
TrainConfigGTAUAV.use_gradcam = False
|
||||
@@ -1,8 +0,0 @@
|
||||
# GTA-UAV Image-heavy: gate initialized high (more image weight).
|
||||
# query = sigma(0.9) * drone + 0.1 * text
|
||||
# Minimal text contribution test.
|
||||
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
TrainConfigGTAUAV.init_gate = 0.9
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/image_heavy"
|
||||
@@ -1,8 +0,0 @@
|
||||
# GTA-UAV Text-heavy: gate initialized low (more text weight).
|
||||
# query = sigma(0.3) * drone + 0.7 * text
|
||||
# Stress test for text contribution.
|
||||
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
TrainConfigGTAUAV.init_gate = 0.3
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/text_heavy"
|
||||
@@ -1,9 +0,0 @@
|
||||
# Text-heavy: gate initialized low (more text weight).
|
||||
# query = sigma(0.3) * drone + 0.7 * text
|
||||
|
||||
include 'balanced.gin'
|
||||
|
||||
DualEncoderCaptionTest.init_gate = 0.3
|
||||
GatedFusion.init_gate = 0.3
|
||||
|
||||
TrainConfig.output_dir = "out/caption_test/text_heavy"
|
||||
12
in/config_files/gtauav_balanced/models.gin
Normal file
12
in/config_files/gtauav_balanced/models.gin
Normal file
@@ -0,0 +1,12 @@
|
||||
# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
|
||||
ModelsCommonConfig.backbone = 'dinov3'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
|
||||
DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
|
||||
DINOv3ModelsConfig.shared_encoder = True
|
||||
DINOv3ModelsConfig.mona_bottleneck = 64
|
||||
DINOv3ModelsConfig.mona_last_n_blocks = 12
|
||||
DINOv3ModelsConfig.lora_rank = 4
|
||||
14
in/config_files/gtauav_balanced/pipeline.gin
Normal file
14
in/config_files/gtauav_balanced/pipeline.gin
Normal file
@@ -0,0 +1,14 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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
|
||||
11
in/config_files/gtauav_balanced_asym/models.gin
Normal file
11
in/config_files/gtauav_balanced_asym/models.gin
Normal file
@@ -0,0 +1,11 @@
|
||||
ModelsCommonConfig.backbone = 'dinov3'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
|
||||
DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
|
||||
DINOv3ModelsConfig.shared_encoder = False
|
||||
DINOv3ModelsConfig.mona_bottleneck = 64
|
||||
DINOv3ModelsConfig.mona_last_n_blocks = 24
|
||||
DINOv3ModelsConfig.lora_rank = 4
|
||||
15
in/config_files/gtauav_balanced_asym/pipeline.gin
Normal file
15
in/config_files/gtauav_balanced_asym/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/balanced_asym'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
91
in/config_files/gtauav_balanced_sofia/models.gin
Normal file
91
in/config_files/gtauav_balanced_sofia/models.gin
Normal file
@@ -0,0 +1,91 @@
|
||||
# SOFIA v7.1 Tiny preset (~5M params) with text fusion (Text-FiLM mid-level
|
||||
# in SAT and UAV heads + GatedFusion late-level on descriptors).
|
||||
#
|
||||
# Tiny-specific notes:
|
||||
# - num_heads_s3/s4 = 4 (channels 176/224 not divisible by 8)
|
||||
# - mamba_headdim = 16 (channels not divisible by default 64)
|
||||
# - mamba_variant = 'mamba1' (Mamba-2 torch fallback bug for these dims)
|
||||
# - d_descriptor = 1024 (override from preset M default 512)
|
||||
# - text fusion enabled (override from preset M default disabled)
|
||||
|
||||
ModelsCommonConfig.backbone = 'sofia_v71'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
# Variant label (informational).
|
||||
SOFIAv71ModelsConfig.variant_label = 'Tiny'
|
||||
|
||||
# Input.
|
||||
SOFIAv71ModelsConfig.input_size = 256
|
||||
SOFIAv71ModelsConfig.in_channels = 3
|
||||
|
||||
# Stem (Tiny dims).
|
||||
SOFIAv71ModelsConfig.stem_mid = 16
|
||||
SOFIAv71ModelsConfig.stem_out = 32
|
||||
|
||||
# Backbone dimensions (Tiny).
|
||||
SOFIAv71ModelsConfig.embed_dims = [48, 96, 176, 224]
|
||||
SOFIAv71ModelsConfig.depths = [2, 3, 4, 2]
|
||||
|
||||
# Stage 1-2 block params (default).
|
||||
SOFIAv71ModelsConfig.mbconv_expand = 4
|
||||
SOFIAv71ModelsConfig.se_ratio = 16
|
||||
SOFIAv71ModelsConfig.strip_kernel_s1 = 7
|
||||
SOFIAv71ModelsConfig.strip_kernel_s2 = 5
|
||||
SOFIAv71ModelsConfig.mix_kernels = [3, 5, 7]
|
||||
SOFIAv71ModelsConfig.use_dcn_strip = True
|
||||
|
||||
# Stage 3-4 (MambaVision). Tiny: mamba1 to bypass torch fallback bug.
|
||||
SOFIAv71ModelsConfig.mamba_d_state = 16
|
||||
SOFIAv71ModelsConfig.mamba_dt_rank = None
|
||||
SOFIAv71ModelsConfig.mamba_backend = 'auto'
|
||||
SOFIAv71ModelsConfig.mamba_variant = 'mamba1'
|
||||
|
||||
# Mamba-2 tunables (used when mamba_variant='mamba2'; Tiny would need
|
||||
# headdim=16 because 176 % 64 != 0 and 224 % 64 != 0).
|
||||
SOFIAv71ModelsConfig.mamba_d_state_mamba2 = 64
|
||||
SOFIAv71ModelsConfig.mamba_headdim = 16
|
||||
SOFIAv71ModelsConfig.mamba_expand = 2
|
||||
SOFIAv71ModelsConfig.mamba_d_conv = 4
|
||||
SOFIAv71ModelsConfig.mamba_n_directions = 2
|
||||
|
||||
# Heads / attention (Tiny: heads=4).
|
||||
SOFIAv71ModelsConfig.num_heads_s3 = 4
|
||||
SOFIAv71ModelsConfig.num_heads_s4 = 4
|
||||
SOFIAv71ModelsConfig.use_strip_branch_s3 = True
|
||||
SOFIAv71ModelsConfig.use_strip_branch_s4 = False
|
||||
SOFIAv71ModelsConfig.ffn_expand = 4
|
||||
|
||||
# EVSS bridge (off by default).
|
||||
SOFIAv71ModelsConfig.use_evss_bridge = False
|
||||
SOFIAv71ModelsConfig.evss_bridge_locations = ['pre_stage3']
|
||||
|
||||
# Neck (Tiny).
|
||||
SOFIAv71ModelsConfig.neck_channels = 128
|
||||
|
||||
# CVGL Head.
|
||||
SOFIAv71ModelsConfig.d_descriptor = 1024
|
||||
SOFIAv71ModelsConfig.use_asymmetric_heads = True
|
||||
SOFIAv71ModelsConfig.chp_rings = 8
|
||||
SOFIAv71ModelsConfig.chp_angles = 16
|
||||
SOFIAv71ModelsConfig.chp_harmonics = 4
|
||||
SOFIAv71ModelsConfig.use_film_altitude = True
|
||||
SOFIAv71ModelsConfig.altitude_norm = 500.0
|
||||
SOFIAv71ModelsConfig.ring_count = 4
|
||||
SOFIAv71ModelsConfig.use_ring_aux = True
|
||||
|
||||
# Text fusion enabled.
|
||||
SOFIAv71ModelsConfig.return_normalized = True
|
||||
SOFIAv71ModelsConfig.use_text_film_sat = True
|
||||
SOFIAv71ModelsConfig.use_text_film_uav = True
|
||||
SOFIAv71ModelsConfig.text_film_dim = 1024
|
||||
SOFIAv71ModelsConfig.text_film_hidden = 256
|
||||
|
||||
# Sharing / KD / deploy.
|
||||
SOFIAv71ModelsConfig.share_stages_1_2 = True
|
||||
SOFIAv71ModelsConfig.enable_kd_taps = True
|
||||
SOFIAv71ModelsConfig.precision = 'fp16'
|
||||
|
||||
# LoRA.
|
||||
SOFIAv71ModelsConfig.lora_rank = 4
|
||||
15
in/config_files/gtauav_balanced_sofia/pipeline.gin
Normal file
15
in/config_files/gtauav_balanced_sofia/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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_sofia'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
30
in/config_files/gtauav_balanced_sofia_v1/models.gin
Normal file
30
in/config_files/gtauav_balanced_sofia_v1/models.gin
Normal file
@@ -0,0 +1,30 @@
|
||||
# SOFIA v1 'tiny' variant (~1M params) with text fusion (Text-FiLM mid-level
|
||||
# in SAT and UAV heads + AltitudeFiLM in UAV head + GatedFusion late-level).
|
||||
|
||||
ModelsCommonConfig.backbone = 'sofia_v1'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
# Backbone.
|
||||
SOFIAv1ModelsConfig.variant_label = 'tiny'
|
||||
SOFIAv1ModelsConfig.in_channels = 3
|
||||
SOFIAv1ModelsConfig.input_size = 256
|
||||
SOFIAv1ModelsConfig.dcn_variant = 'v2'
|
||||
|
||||
# Heads.
|
||||
SOFIAv1ModelsConfig.d_descriptor = 1024
|
||||
SOFIAv1ModelsConfig.return_normalized = False
|
||||
|
||||
# Altitude-FiLM.
|
||||
SOFIAv1ModelsConfig.use_film_altitude = True
|
||||
SOFIAv1ModelsConfig.altitude_norm = 500.0
|
||||
|
||||
# Text-FiLM.
|
||||
SOFIAv1ModelsConfig.use_text_film_uav = True
|
||||
SOFIAv1ModelsConfig.use_text_film_sat = True
|
||||
SOFIAv1ModelsConfig.text_film_dim = 1024
|
||||
SOFIAv1ModelsConfig.text_film_hidden = 256
|
||||
|
||||
# LoRA on DGTRS-CLIP.
|
||||
SOFIAv1ModelsConfig.lora_rank = 4
|
||||
15
in/config_files/gtauav_balanced_sofia_v1/pipeline.gin
Normal file
15
in/config_files/gtauav_balanced_sofia_v1/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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_sofia_v1'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
11
in/config_files/gtauav_balanced_stripnet/models.gin
Normal file
11
in/config_files/gtauav_balanced_stripnet/models.gin
Normal file
@@ -0,0 +1,11 @@
|
||||
# StripNet small backbone (frozen) + Conv-MONA on last 2 stages.
|
||||
ModelsCommonConfig.backbone = 'stripnet'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
|
||||
StripNetModelsConfig.stripnet_freeze = True
|
||||
StripNetModelsConfig.stripnet_mona_last_n_stages = 2
|
||||
StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
|
||||
StripNetModelsConfig.lora_rank = 4
|
||||
15
in/config_files/gtauav_balanced_stripnet/pipeline.gin
Normal file
15
in/config_files/gtauav_balanced_stripnet/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/balanced_stripnet'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
11
in/config_files/gtauav_balanced_stripnet_unfrozen/models.gin
Normal file
11
in/config_files/gtauav_balanced_stripnet_unfrozen/models.gin
Normal file
@@ -0,0 +1,11 @@
|
||||
# StripNet small backbone (frozen) + Conv-MONA on last 2 stages.
|
||||
ModelsCommonConfig.backbone = 'stripnet'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
|
||||
StripNetModelsConfig.stripnet_freeze = False
|
||||
StripNetModelsConfig.stripnet_mona_last_n_stages = 2
|
||||
StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
|
||||
StripNetModelsConfig.lora_rank = 4
|
||||
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/balanced_stripnet_unfrozen'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
12
in/config_files/gtauav_baseline/models.gin
Normal file
12
in/config_files/gtauav_baseline/models.gin
Normal file
@@ -0,0 +1,12 @@
|
||||
# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
|
||||
ModelsCommonConfig.backbone = 'dinov3'
|
||||
ModelsCommonConfig.baseline_mode = True
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
|
||||
DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
|
||||
DINOv3ModelsConfig.shared_encoder = True
|
||||
DINOv3ModelsConfig.mona_bottleneck = 64
|
||||
DINOv3ModelsConfig.mona_last_n_blocks = 12
|
||||
DINOv3ModelsConfig.lora_rank = 4
|
||||
14
in/config_files/gtauav_baseline/pipeline.gin
Normal file
14
in/config_files/gtauav_baseline/pipeline.gin
Normal file
@@ -0,0 +1,14 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/baseline_inbatch'
|
||||
PipelineConfig.resume_from = None
|
||||
12
in/config_files/gtauav_baseline_asym/models.gin
Normal file
12
in/config_files/gtauav_baseline_asym/models.gin
Normal file
@@ -0,0 +1,12 @@
|
||||
# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
|
||||
ModelsCommonConfig.backbone = 'dinov3'
|
||||
ModelsCommonConfig.baseline_mode = True
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
|
||||
DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
|
||||
DINOv3ModelsConfig.shared_encoder = False
|
||||
DINOv3ModelsConfig.mona_bottleneck = 64
|
||||
DINOv3ModelsConfig.mona_last_n_blocks = 24
|
||||
DINOv3ModelsConfig.lora_rank = 4
|
||||
14
in/config_files/gtauav_baseline_asym/pipeline.gin
Normal file
14
in/config_files/gtauav_baseline_asym/pipeline.gin
Normal file
@@ -0,0 +1,14 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/baseline_asym'
|
||||
PipelineConfig.resume_from = None
|
||||
92
in/config_files/gtauav_baseline_sofia/models.gin
Normal file
92
in/config_files/gtauav_baseline_sofia/models.gin
Normal file
@@ -0,0 +1,92 @@
|
||||
# SOFIA v7.1 Tiny preset (~5M params) with text fusion (Text-FiLM mid-level
|
||||
# in SAT and UAV heads + GatedFusion late-level on descriptors).
|
||||
#
|
||||
# Tiny-specific notes:
|
||||
# - num_heads_s3/s4 = 4 (channels 176/224 not divisible by 8)
|
||||
# - mamba_headdim = 16 (channels not divisible by default 64)
|
||||
# - mamba_variant = 'mamba1' (Mamba-2 torch fallback bug for these dims)
|
||||
# - d_descriptor = 1024 (override from preset M default 512)
|
||||
# - text fusion enabled (override from preset M default disabled)
|
||||
|
||||
ModelsCommonConfig.backbone = 'sofia_v71'
|
||||
ModelsCommonConfig.baseline_mode = True
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
# Variant label (informational).
|
||||
SOFIAv71ModelsConfig.variant_label = 'Tiny'
|
||||
|
||||
# Input.
|
||||
SOFIAv71ModelsConfig.input_size = 256
|
||||
SOFIAv71ModelsConfig.in_channels = 3
|
||||
|
||||
# Stem (Tiny dims).
|
||||
SOFIAv71ModelsConfig.stem_mid = 16
|
||||
SOFIAv71ModelsConfig.stem_out = 32
|
||||
|
||||
# Backbone dimensions (Tiny).
|
||||
SOFIAv71ModelsConfig.embed_dims = [48, 96, 176, 224]
|
||||
SOFIAv71ModelsConfig.depths = [2, 3, 4, 2]
|
||||
|
||||
# Stage 1-2 block params (default).
|
||||
SOFIAv71ModelsConfig.mbconv_expand = 4
|
||||
SOFIAv71ModelsConfig.se_ratio = 16
|
||||
SOFIAv71ModelsConfig.strip_kernel_s1 = 7
|
||||
SOFIAv71ModelsConfig.strip_kernel_s2 = 5
|
||||
SOFIAv71ModelsConfig.mix_kernels = [3, 5, 7]
|
||||
SOFIAv71ModelsConfig.use_dcn_strip = True
|
||||
|
||||
# Stage 3-4 (MambaVision). Tiny: mamba1 to bypass torch fallback bug.
|
||||
SOFIAv71ModelsConfig.mamba_d_state = 16
|
||||
SOFIAv71ModelsConfig.mamba_dt_rank = None
|
||||
SOFIAv71ModelsConfig.mamba_backend = 'auto'
|
||||
SOFIAv71ModelsConfig.mamba_variant = 'mamba1'
|
||||
|
||||
# Mamba-2 tunables (used when mamba_variant='mamba2'; Tiny would need
|
||||
# headdim=16 because 176 % 64 != 0 and 224 % 64 != 0).
|
||||
SOFIAv71ModelsConfig.mamba_d_state_mamba2 = 64
|
||||
SOFIAv71ModelsConfig.mamba_headdim = 16
|
||||
SOFIAv71ModelsConfig.mamba_expand = 2
|
||||
SOFIAv71ModelsConfig.mamba_d_conv = 4
|
||||
SOFIAv71ModelsConfig.mamba_n_directions = 2
|
||||
|
||||
# Heads / attention (Tiny: heads=4).
|
||||
SOFIAv71ModelsConfig.num_heads_s3 = 4
|
||||
SOFIAv71ModelsConfig.num_heads_s4 = 4
|
||||
SOFIAv71ModelsConfig.use_strip_branch_s3 = True
|
||||
SOFIAv71ModelsConfig.use_strip_branch_s4 = False
|
||||
SOFIAv71ModelsConfig.ffn_expand = 4
|
||||
|
||||
# EVSS bridge (off by default).
|
||||
SOFIAv71ModelsConfig.use_evss_bridge = False
|
||||
SOFIAv71ModelsConfig.evss_bridge_locations = ['pre_stage3']
|
||||
|
||||
# Neck (Tiny).
|
||||
SOFIAv71ModelsConfig.neck_channels = 128
|
||||
|
||||
# CVGL Head.
|
||||
SOFIAv71ModelsConfig.d_descriptor = 1024
|
||||
SOFIAv71ModelsConfig.use_asymmetric_heads = True
|
||||
SOFIAv71ModelsConfig.chp_rings = 8
|
||||
SOFIAv71ModelsConfig.chp_angles = 16
|
||||
SOFIAv71ModelsConfig.chp_harmonics = 4
|
||||
SOFIAv71ModelsConfig.use_film_altitude = True
|
||||
SOFIAv71ModelsConfig.altitude_norm = 500.0
|
||||
SOFIAv71ModelsConfig.ring_count = 4
|
||||
SOFIAv71ModelsConfig.use_ring_aux = True
|
||||
|
||||
# Text fusion enabled.
|
||||
SOFIAv71ModelsConfig.return_normalized = True
|
||||
SOFIAv71ModelsConfig.use_text_film_sat = True
|
||||
SOFIAv71ModelsConfig.use_text_film_uav = True
|
||||
SOFIAv71ModelsConfig.text_film_dim = 1024
|
||||
SOFIAv71ModelsConfig.text_film_hidden = 256
|
||||
|
||||
# Sharing / KD / deploy.
|
||||
SOFIAv71ModelsConfig.share_stages_1_2 = True
|
||||
SOFIAv71ModelsConfig.enable_kd_taps = True
|
||||
SOFIAv71ModelsConfig.precision = 'fp16'
|
||||
|
||||
# LoRA.
|
||||
SOFIAv71ModelsConfig.lora_rank = 4
|
||||
|
||||
15
in/config_files/gtauav_baseline_sofia/pipeline.gin
Normal file
15
in/config_files/gtauav_baseline_sofia/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/baseline_sofia'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
31
in/config_files/gtauav_baseline_sofia_v1/models.gin
Normal file
31
in/config_files/gtauav_baseline_sofia_v1/models.gin
Normal file
@@ -0,0 +1,31 @@
|
||||
# SOFIA v1 'tiny' variant (~1M params) with text fusion (Text-FiLM mid-level
|
||||
# in SAT and UAV heads + AltitudeFiLM in UAV head + GatedFusion late-level).
|
||||
|
||||
ModelsCommonConfig.backbone = 'sofia_v1'
|
||||
ModelsCommonConfig.baseline_mode = True
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
# Backbone.
|
||||
SOFIAv1ModelsConfig.variant_label = 'tiny'
|
||||
SOFIAv1ModelsConfig.in_channels = 3
|
||||
SOFIAv1ModelsConfig.input_size = 256
|
||||
SOFIAv1ModelsConfig.dcn_variant = 'v2'
|
||||
|
||||
# Heads.
|
||||
SOFIAv1ModelsConfig.d_descriptor = 1024
|
||||
SOFIAv1ModelsConfig.return_normalized = False
|
||||
|
||||
# Altitude-FiLM.
|
||||
SOFIAv1ModelsConfig.use_film_altitude = True
|
||||
SOFIAv1ModelsConfig.altitude_norm = 500.0
|
||||
|
||||
# Text-FiLM.
|
||||
SOFIAv1ModelsConfig.use_text_film_uav = True
|
||||
SOFIAv1ModelsConfig.use_text_film_sat = True
|
||||
SOFIAv1ModelsConfig.text_film_dim = 1024
|
||||
SOFIAv1ModelsConfig.text_film_hidden = 256
|
||||
|
||||
# LoRA on DGTRS-CLIP.
|
||||
SOFIAv1ModelsConfig.lora_rank = 4
|
||||
|
||||
15
in/config_files/gtauav_baseline_sofia_v1/pipeline.gin
Normal file
15
in/config_files/gtauav_baseline_sofia_v1/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/baseline_sofia_v1'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
12
in/config_files/gtauav_baseline_stripnet/models.gin
Normal file
12
in/config_files/gtauav_baseline_stripnet/models.gin
Normal file
@@ -0,0 +1,12 @@
|
||||
# StripNet small backbone (frozen) + Conv-MONA on last 2 stages.
|
||||
ModelsCommonConfig.backbone = 'stripnet'
|
||||
ModelsCommonConfig.baseline_mode = True
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
|
||||
StripNetModelsConfig.stripnet_freeze = True
|
||||
StripNetModelsConfig.stripnet_mona_last_n_stages = 2
|
||||
StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
|
||||
StripNetModelsConfig.lora_rank = 4
|
||||
|
||||
15
in/config_files/gtauav_baseline_stripnet/pipeline.gin
Normal file
15
in/config_files/gtauav_baseline_stripnet/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/baseline_stripnet'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
11
in/config_files/gtauav_baseline_stripnet_unfrozen/models.gin
Normal file
11
in/config_files/gtauav_baseline_stripnet_unfrozen/models.gin
Normal file
@@ -0,0 +1,11 @@
|
||||
# StripNet small backbone (unfrozen) + Conv-MONA on last 2 stages.
|
||||
ModelsCommonConfig.backbone = 'stripnet'
|
||||
ModelsCommonConfig.baseline_mode = True
|
||||
ModelsCommonConfig.init_gate = 0.7
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
StripNetModelsConfig.stripnet_path = 'nn_models/STRIPNET/stripnet_s.pth'
|
||||
StripNetModelsConfig.stripnet_freeze = False
|
||||
StripNetModelsConfig.stripnet_mona_last_n_stages = 2
|
||||
StripNetModelsConfig.stripnet_backbone_lr_factor = 0.1
|
||||
StripNetModelsConfig.lora_rank = 4
|
||||
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/baseline_stripnet_unfrozen'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
13
in/config_files/gtauav_image_heavy/models.gin
Normal file
13
in/config_files/gtauav_image_heavy/models.gin
Normal file
@@ -0,0 +1,13 @@
|
||||
# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
|
||||
ModelsCommonConfig.backbone = 'dinov3'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.9
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
|
||||
DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
|
||||
DINOv3ModelsConfig.shared_encoder = True
|
||||
DINOv3ModelsConfig.mona_bottleneck = 64
|
||||
DINOv3ModelsConfig.mona_last_n_blocks = 12
|
||||
DINOv3ModelsConfig.lora_rank = 4
|
||||
|
||||
15
in/config_files/gtauav_image_heavy/pipeline.gin
Normal file
15
in/config_files/gtauav_image_heavy/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/image_heavy'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
13
in/config_files/gtauav_text_heavy/models.gin
Normal file
13
in/config_files/gtauav_text_heavy/models.gin
Normal file
@@ -0,0 +1,13 @@
|
||||
# DINOv3 shared encoder + MONA-12 + DGTRS-CLIP with text.
|
||||
ModelsCommonConfig.backbone = 'dinov3'
|
||||
ModelsCommonConfig.baseline_mode = False
|
||||
ModelsCommonConfig.init_gate = 0.3
|
||||
ModelsCommonConfig.lrsclip_path = 'nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt'
|
||||
|
||||
DINOv3ModelsConfig.dino_web_path = 'nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth'
|
||||
DINOv3ModelsConfig.dino_sat_path = 'nn_models/DINO_SAT/model.safetensors'
|
||||
DINOv3ModelsConfig.shared_encoder = True
|
||||
DINOv3ModelsConfig.mona_bottleneck = 64
|
||||
DINOv3ModelsConfig.mona_last_n_blocks = 12
|
||||
DINOv3ModelsConfig.lora_rank = 4
|
||||
|
||||
15
in/config_files/gtauav_text_heavy/pipeline.gin
Normal file
15
in/config_files/gtauav_text_heavy/pipeline.gin
Normal file
@@ -0,0 +1,15 @@
|
||||
# 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 = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PipelineConfig.caption_root = '/media/servml/SSD_2_2TB/datasets/cvgl_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/text_heavy'
|
||||
PipelineConfig.resume_from = None
|
||||
|
||||
9
in/config_files/hardware_default.gin
Normal file
9
in/config_files/hardware_default.gin
Normal file
@@ -0,0 +1,9 @@
|
||||
# RTX 4090 profile, shared encoder (DINOv3).
|
||||
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
|
||||
HardwareConfig.total_vram_gb = 24.0
|
||||
10
in/config_files/hardware_no_gradckpt.gin
Normal file
10
in/config_files/hardware_no_gradckpt.gin
Normal file
@@ -0,0 +1,10 @@
|
||||
# SOFIA v7.1 from-scratch — keep activations live (no gradient checkpointing).
|
||||
HardwareConfig.device = 'cuda'
|
||||
HardwareConfig.batch_size = 8
|
||||
HardwareConfig.grad_accum_steps = 8
|
||||
HardwareConfig.num_workers = 4
|
||||
HardwareConfig.use_amp = True
|
||||
HardwareConfig.gradient_checkpointing = False
|
||||
HardwareConfig.reserve_gb = 2.0
|
||||
HardwareConfig.total_vram_gb = 24.0
|
||||
|
||||
20
in/config_files/preprocess/preprocess.gin
Normal file
20
in/config_files/preprocess/preprocess.gin
Normal file
@@ -0,0 +1,20 @@
|
||||
# Preprocessing config used by scripts/make_split.py and
|
||||
# scripts/filter_segmentation.py. Independent from training pipeline.
|
||||
|
||||
# Inputs.
|
||||
PreprocessConfig.rgb_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR'
|
||||
PreprocessConfig.segm_root = '/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm'
|
||||
|
||||
# make_split.py — 80/20 split with seed=42.
|
||||
PreprocessConfig.split_ratio = 0.8
|
||||
PreprocessConfig.split_seed = 42
|
||||
PreprocessConfig.split_input_train = 'cross-area-drone2sate-train.json'
|
||||
PreprocessConfig.split_input_test = 'cross-area-drone2sate-test.json'
|
||||
PreprocessConfig.split_output_dir = 'meta'
|
||||
PreprocessConfig.split_output_train = 'train_80.json'
|
||||
PreprocessConfig.split_output_test = 'test_20.json'
|
||||
|
||||
# filter_segmentation.py — exclude images with >=90% background+water.
|
||||
PreprocessConfig.seg_threshold = 0.90
|
||||
PreprocessConfig.seg_exclude_classes = [0, 4]
|
||||
PreprocessConfig.seg_filter_output = 'meta/seg_filter.json'
|
||||
16
in/config_files/tracking.gin
Normal file
16
in/config_files/tracking.gin
Normal file
@@ -0,0 +1,16 @@
|
||||
# Default tracking: TensorBoard on, W&B off, no Grad-CAM, no profiler.
|
||||
TrackingConfig.use_wandb = False
|
||||
TrackingConfig.use_tb = True
|
||||
TrackingConfig.wandb_project = 'caption-test-gtauav'
|
||||
TrackingConfig.wandb_run_name = None
|
||||
TrackingConfig.wandb_entity = None
|
||||
|
||||
TrackingConfig.log_grad_norms = True
|
||||
TrackingConfig.use_gradcam = False
|
||||
TrackingConfig.gradcam_every = 5
|
||||
TrackingConfig.gradcam_samples = 8
|
||||
|
||||
TrackingConfig.use_profiler = False
|
||||
TrackingConfig.profiler_warmup = 3
|
||||
TrackingConfig.profiler_active = 5
|
||||
|
||||
31
in/config_files/training.gin
Normal file
31
in/config_files/training.gin
Normal file
@@ -0,0 +1,31 @@
|
||||
# Loss + optimizer + sampler — symmetric InfoNCE, AdamW, mutex sampler.
|
||||
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.tau_final = 0.01
|
||||
TrainingConfig.weight_q2g = 0.6
|
||||
TrainingConfig.weight_g2q = 0.4
|
||||
TrainingConfig.hard_mining_k = 0
|
||||
TrainingConfig.neg_bank_size = 0
|
||||
|
||||
# WeightedInfoNCE-only (unused when loss_type='symmetric').
|
||||
TrainingConfig.weighted_loss_k = 5.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
|
||||
TrainingConfig.dss_lsh_num_tables = 8
|
||||
TrainingConfig.dss_lsh_num_bits = 14
|
||||
TrainingConfig.dss_cache_dir = None
|
||||
TrainingConfig.use_mutex_sampler = True
|
||||
297
pre_refactor_analysis.md
Normal file
297
pre_refactor_analysis.md
Normal file
@@ -0,0 +1,297 @@
|
||||
# Шаг 1 — Что не соответствует стандарту в текущем состоянии
|
||||
|
||||
> **Только диагностика.** Никаких решений по рефакторингу — это пойдёт на следующих шагах. Здесь только: что есть сейчас, какое правило стандарта нарушено, и каков масштаб проблемы.
|
||||
|
||||
---
|
||||
|
||||
## Текущее состояние
|
||||
|
||||
| Что | Где | Сколько |
|
||||
|---|---|---|
|
||||
| Гин-конфиги (взаимовключаемые через `include`) | `conf/` | 12 файлов: `balanced.gin`, `baseline_no_text.gin`, `text_heavy.gin` (v2 UAV-GeoLoc), `gtauav_balanced.gin`, `gtauav_baseline.gin`, `gtauav_balanced_asym.gin`, `gtauav_baseline_asym.gin`, `gtauav_balanced_stripnet.gin`, `gtauav_balanced_stripnet_unfrozen.gin`, `gtauav_baseline_stripnet.gin`, `gtauav_baseline_stripnet_unfrozen.gin`, `gtauav_text_heavy.gin`, `gtauav_image_heavy.gin` (по скриншоту в локальной копии есть ещё 4 sofia-варианта) |
|
||||
| Главный `@gin.configurable + @dataclass` | `src/training/train_gtauav.py::TrainConfigGTAUAV` | 1 класс, 50+ полей |
|
||||
| Legacy `@gin.configurable` на классе тренировки | `src/training/train.py::TrainConfig` (v2) | 1 класс, 14 полей |
|
||||
| `@gin.configurable` на не-конфиг классах | `src/losses/multi_infonce.py::InfoNCELoss`, `src/losses/weighted_infonce.py::WeightedInfoNCELoss`, `src/datasets/visloc_with_captions.py::GeoLocCaptionDataset` | 3 класса |
|
||||
| Локальный `@dataclass` для модели (без gin) | `src/models/sofia_v1/config.py::SOFIAv1Config` | 1 класс |
|
||||
| Module-level пути | `src/training/train_gtauav.py`, `src/datasets/gtauav_dataset.py`, `scripts/make_split.py`, `scripts/filter_segmentation.py` | 4 файла |
|
||||
| `argparse` | `src/training/train_gtauav.py`, `src/training/train.py`, `scripts/make_split.py`, `scripts/filter_segmentation.py` | 4 файла |
|
||||
|
||||
---
|
||||
|
||||
## Несоответствия стандарту
|
||||
|
||||
### 1. `@gin.configurable + @dataclass` — критическое нарушение
|
||||
|
||||
**Правило** (`Стандарт_написания_кода_для_DL_CV.md` §3.1, `Reference Examples → Anti-patterns`):
|
||||
|
||||
> «**Запрещено** использовать `dataclass` совместно с gin»
|
||||
|
||||
**Что нарушает:**
|
||||
|
||||
`src/training/train_gtauav.py`:
|
||||
```python
|
||||
@gin.configurable(module="src.training.train_gtauav")
|
||||
@dataclass # ← FORBIDDEN
|
||||
class TrainConfigGTAUAV:
|
||||
train_json: str = _TRAIN_JSON
|
||||
test_json: str = _TEST_JSON
|
||||
rgb_root: str = _RGB_ROOT
|
||||
# ... 50+ полей всего:
|
||||
# пути к данным, к моделям, training schedule, model arch (mona, gates,
|
||||
# baseline, stripnet, asymmetric), loss params, sampler params,
|
||||
# tracking flags (wandb/tb/gradcam/profiler)
|
||||
```
|
||||
|
||||
**Почему это нарушение:** `@dataclass` авто-генерирует `__init__`, в который gin **тоже** инъектирует параметры. Получается двойная магия: dataclass читает type hints для генерации сигнатуры, gin читает гин-биндинги для подмены значений. На практике работает, но это именно та комбинация, которую стандарт запрещает.
|
||||
|
||||
**Масштаб:** **1 класс**, но он держит **всё** — путь, hardware, model, loss, sampler, tracking. Это блокирует разделение конфига на оси.
|
||||
|
||||
---
|
||||
|
||||
### 2. `@gin.configurable` на не-конфиг классах — критическое нарушение
|
||||
|
||||
**Правило** (стандарт §3.1):
|
||||
|
||||
> «Декоратор `@gin.configurable` **только на классах** [конфигурации, *не на бизнес-логике*]»
|
||||
|
||||
**Что нарушает:**
|
||||
|
||||
| Файл | Класс | Параметры (которые проникают в `.gin`) |
|
||||
|---|---|---|
|
||||
| `src/losses/multi_infonce.py` | `InfoNCELoss` | `temperature_init`, `temperature_final`, `label_smoothing`, `weight_q2g`, `weight_g2q`, `learnable_temperature`, `tau_min`, `tau_max`, `hard_mining_k` |
|
||||
| `src/losses/weighted_infonce.py` | `WeightedInfoNCELoss` | `temperature_init`, `learnable_temperature`, `label_smoothing`, `k`, `tau_min`, `tau_max` |
|
||||
| `src/datasets/visloc_with_captions.py` (v2) | `GeoLocCaptionDataset` | `query_file`, `data_root`, `image_transform`, `drop_caption_prob`, `seed` |
|
||||
|
||||
**Почему это нарушение:** конфиг и бизнес-логика — разные слои. Когда `nn.Module` или `Dataset` декорированы `@gin.configurable`, гин лезет в **их** `__init__` помимо лезения в `TrainConfigGTAUAV.__init__`. В `gtauav_balanced.gin` это видно прямо:
|
||||
|
||||
```gin
|
||||
# Параметры дублируются — раз в TrainConfigGTAUAV, раз в InfoNCELoss:
|
||||
TrainConfigGTAUAV.tau_init = 0.07
|
||||
TrainConfigGTAUAV.label_smoothing = 0.1
|
||||
TrainConfigGTAUAV.weight_q2g = 0.6
|
||||
TrainConfigGTAUAV.weight_g2q = 0.4
|
||||
|
||||
# ---- InfoNCELoss (gin-configurable) ----
|
||||
InfoNCELoss.temperature_init = 0.07 # ← дубль tau_init
|
||||
InfoNCELoss.label_smoothing = 0.1 # ← дубль
|
||||
InfoNCELoss.weight_q2g = 0.6 # ← дубль
|
||||
InfoNCELoss.weight_g2q = 0.4 # ← дубль
|
||||
InfoNCELoss.tau_min = 0.01
|
||||
InfoNCELoss.tau_max = 0.1
|
||||
InfoNCELoss.hard_mining_k = 0
|
||||
```
|
||||
|
||||
**Это активный источник тихих багов:** если кто-то поменяет `TrainConfigGTAUAV.tau_init = 0.05`, а `InfoNCELoss.temperature_init` забудет — обучение пойдёт с `0.05` в логике trainer-а и `0.07` в самой loss-функции. Никаких ошибок не будет, метрики просто будут странными.
|
||||
|
||||
**Масштаб:** **3 класса**, но самый болезненный — `InfoNCELoss`, потому что у него **самая большая зона перекрытия** с `TrainConfigGTAUAV`.
|
||||
|
||||
---
|
||||
|
||||
### 3. `argparse` — критическое нарушение
|
||||
|
||||
**Правило** (стандарт §3.4):
|
||||
|
||||
> «**Запрещён argparse** — все параметры из .gin файлов»
|
||||
|
||||
**Что нарушает:**
|
||||
|
||||
| Файл | Кол-во CLI флагов | Самые проблемные |
|
||||
|---|---|---|
|
||||
| `src/training/train_gtauav.py::main` | ~15 | `--config`, `--baseline`, `--batch-size`, `--epochs`, `--filter-meta`, `--wandb`, `--gradcam`, `--profile`, `--gin-param`, `--resume`, `--output-dir` |
|
||||
| `src/training/train.py::main` (v2) | 1 | `--config` |
|
||||
| `scripts/make_split.py::main` | 3 | `--ratio`, `--seed`, `--output-dir` |
|
||||
| `scripts/filter_segmentation.py::main` | 3 | `--segm-root`, `--threshold`, `--output` |
|
||||
|
||||
**Самый ядовитый паттерн** в `train_gtauav.py`:
|
||||
|
||||
```python
|
||||
parser.add_argument("--gin-param", nargs="*", help="Override gin params from CLI")
|
||||
# ...
|
||||
gin.parse_config_files_and_bindings([cfg_file], extra_bindings)
|
||||
```
|
||||
|
||||
CLI **перекрывает** gin-биндинги. Это создаёт **3 источника правды** на один и тот же параметр: дефолт в `__init__`, значение в `.gin`, значение в `--gin-param`. Какое из них применилось в конкретном запуске — невозможно установить иначе как чтением логов.
|
||||
|
||||
**Масштаб:** **тренировочный код** (`train_gtauav.py`) — критично; **скрипты** (`make_split.py`, `filter_segmentation.py`) — пограничный случай (см. ниже §7).
|
||||
|
||||
---
|
||||
|
||||
### 4. `include` для композиции `.gin` — пограничный случай
|
||||
|
||||
**Правило** (`REQUIREMENTS_GIN_STYLE.md` §8):
|
||||
|
||||
> «Использовать **минимальный** набор возможностей gin: `@gin.configurable`, `gin.parse_config_file()`. Нет: `gin.register()`, `gin.constant()`, `gin.query_parameter()`, **макросы, ссылки между конфигами**»
|
||||
|
||||
**Что есть сейчас:**
|
||||
|
||||
```gin
|
||||
# conf/gtauav_balanced_asym.gin
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
TrainConfigGTAUAV.shared_encoder = False
|
||||
TrainConfigGTAUAV.mona_last_n_blocks = 24
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/balanced_asym"
|
||||
```
|
||||
|
||||
`include` — это **встроенная фича gin** для композиции. Формально она:
|
||||
- ✅ Не входит в явный список запрещённого (`gin.constant`, `gin.register`, `gin.query_parameter`, макросы)
|
||||
- ❌ Попадает под формулировку «**ссылки между конфигами**» в `REQUIREMENTS_GIN_STYLE.md` §8
|
||||
|
||||
**По букве правил это нарушение.** По духу — `include` решает реальную проблему (DRY), но создаёт неявную зависимость: чтобы понять, с какими параметрами запускается `gtauav_balanced_stripnet_unfrozen.gin`, нужно прочитать **3 файла** (`unfrozen` → `stripnet` → `balanced`).
|
||||
|
||||
**Масштаб:** все 14 `gtauav_*.gin` (кроме `gtauav_balanced.gin`, который сам — корень дерева) и v2 (`baseline_no_text.gin`, `text_heavy.gin` тянут `balanced.gin`).
|
||||
|
||||
---
|
||||
|
||||
### 5. Один мега-конфиг на всё — не «жёсткое» нарушение, но против духа
|
||||
|
||||
**Правило** (стандарт §3.3, `Рекомендуемые_gin-config_категории.md`):
|
||||
|
||||
> «Каждый `.gin` → один конфиг-класс»
|
||||
>
|
||||
> «Принцип разделения: если два параметра меняются **вместе** — в одном конфиге. Если **независимо** — в разных»
|
||||
|
||||
**Что есть сейчас:**
|
||||
|
||||
`TrainConfigGTAUAV` миксует **несколько независимых осей изменчивости** в один класс:
|
||||
- Пути к данным/моделям (меняются при смене машины) ↔ training schedule (меняются при смене эксперимента) ↔ model arch (`baseline_mode`, `shared_encoder`, `mona_*`, `stripnet_*`) ↔ loss (`tau_init`, `label_smoothing`, `weight_q2g`, ...) ↔ sampler (`sampler_type`, `dss_*`) ↔ tracking (`use_wandb`, `use_tb`, `use_gradcam`, `use_profiler`)
|
||||
|
||||
В одном `.gin` лежат биндинги для всех этих осей. Когда нужно «то же обучение, но без wandb» — приходится копировать целый `.gin` файл.
|
||||
|
||||
**Это не нарушение явного запрета, но прямое следствие нарушения #1:** пока есть один большой `@dataclass + @gin.configurable`, иначе расположить параметры просто негде.
|
||||
|
||||
**Масштаб:** проникает во все 12 `gtauav_*.gin` файлов одинаково.
|
||||
|
||||
---
|
||||
|
||||
### 6. Module-level хардкод путей — нарушение
|
||||
|
||||
**Правило** (стандарт §6 чеклист):
|
||||
|
||||
> «Нет захардкоженных model ID / промптов / размеров?»
|
||||
|
||||
**Что нарушает:**
|
||||
|
||||
```python
|
||||
# src/training/train_gtauav.py (module level)
|
||||
_RGB_ROOT = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR"
|
||||
_CAPTION_ROOT = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions"
|
||||
_TRAIN_JSON = "meta/train_80.json"
|
||||
_TEST_JSON = "meta/test_20.json"
|
||||
_DINO_WEB = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth"
|
||||
_DINO_SAT = "nn_models/DINO_SAT/model.safetensors"
|
||||
_LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt"
|
||||
|
||||
# src/datasets/gtauav_dataset.py (module level) — ДУБЛЬ:
|
||||
_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
|
||||
_CAPTION_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions")
|
||||
|
||||
# scripts/make_split.py (module level):
|
||||
_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
|
||||
|
||||
# scripts/filter_segmentation.py (module level):
|
||||
SEGM_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm")
|
||||
```
|
||||
|
||||
**Главная проблема — не сам факт хардкода**, а то, что один и тот же путь `/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR` дублируется в **трёх** местах: train_gtauav.py + gtauav_dataset.py + make_split.py. Если переехать на другую машину — нужно править 3 файла; если забыть один — silent breakage.
|
||||
|
||||
**Масштаб:** 4 файла, ~10 module-level констант пути.
|
||||
|
||||
---
|
||||
|
||||
### 7. Скрипты `make_split.py` / `filter_segmentation.py` — спорный случай
|
||||
|
||||
**Что есть:**
|
||||
- `argparse` (3 параметра в каждом)
|
||||
- module-level пути (`_RGB_ROOT`, `SEGM_ROOT`)
|
||||
- Запускаются однократно перед тренировкой
|
||||
- Не интегрированы в `Trainer.run()`
|
||||
|
||||
**Применять ли стандарт «нет argparse» к ним?**
|
||||
|
||||
Стандарт направлен против `argparse` в **тренировочном коде**, где параметров 50+ и нужна воспроизводимость. У `make_split.py` 3 параметра, и это **разовая препроцессинг-утилита**. Здесь буква и дух стандарта расходятся.
|
||||
|
||||
**Это не «несоответствие, требующее немедленной правки»**, а **открытый вопрос для следующего шага**: применять ли gin-стиль к препроцессинг-скриптам или оставить argparse.
|
||||
|
||||
---
|
||||
|
||||
### 8. Sofia models — `@dataclass` без gin (формально не нарушение)
|
||||
|
||||
**Что есть:**
|
||||
|
||||
```python
|
||||
# src/models/sofia_v1/config.py
|
||||
@dataclass
|
||||
class SOFIAv1Config:
|
||||
variant: Literal["tiny_tiny", "tiny", "small", "small_v2"] = "small"
|
||||
in_channels: int = 3
|
||||
input_size: int = 256
|
||||
dcn_variant: Literal["v2", "v4"] = "v2"
|
||||
d_descriptor: int = 1024
|
||||
use_film_altitude: bool = True
|
||||
altitude_norm: float = 500.0
|
||||
use_text_film_uav: bool = True
|
||||
use_text_film_sat: bool = True
|
||||
text_film_dim: int = 1024
|
||||
text_film_hidden: int = 256
|
||||
```
|
||||
|
||||
**Стандарт** запрещает **`@gin.configurable + @dataclass`**, но не запрещает `@dataclass` сам по себе. `SOFIAv1Config` без gin — формально стандарт **не нарушает**.
|
||||
|
||||
**Однако:** в `gtauav_balanced_sofia*.gin` (по скриншоту локально есть, в репо пока нет) параметры sofia, очевидно, прокидываются в `TrainConfigGTAUAV` как обычные поля. То есть sofia сейчас живёт в **двух** местах:
|
||||
- `SOFIAv1Config` (dataclass, в коде модели)
|
||||
- `TrainConfigGTAUAV.sofia_*` (если там такие есть) или только через позиционное создание `SOFIAv1Config()` в `Trainer`
|
||||
|
||||
Это не нарушение, но **архитектурный раскол**: dataclass-конфиги внутри модельных подсистем + gin-конфиги снаружи.
|
||||
|
||||
**Это тоже открытый вопрос для следующего шага**, не текущая проблема.
|
||||
|
||||
---
|
||||
|
||||
### 9. Прочее по чеклисту стандарта (мелочи)
|
||||
|
||||
| Пункт | Состояние | Файлы |
|
||||
|---|---|---|
|
||||
| `from __future__ import annotations` первой строкой | ✅ есть в ключевых файлах | проверено в `train_gtauav.py`, `gtauav_dataset.py`, `multi_infonce.py`, `weighted_infonce.py`, `train.py` |
|
||||
| Строгие type hints | ✅ в основном | пара мест с `dict` без параметров (`_atomic_save(obj: dict)`) |
|
||||
| Google-style docstrings | ✅ есть, качество хорошее | — |
|
||||
| `@torch.inference_mode()` вместо `@torch.no_grad()` | ❌ используется `@torch.no_grad()` | `train_gtauav.py::_evaluate`, `_embed_drone_queries` |
|
||||
| Atomic writes | ⚠️ есть, но без cleanup на ошибке | `train_gtauav.py::_atomic_save` (нет `try/except` — `.tmp` остаётся при сбое) |
|
||||
| Английский в коде/комментариях | ✅ есть | — |
|
||||
| Импорты stdlib → third-party → local | ✅ есть | — |
|
||||
|
||||
---
|
||||
|
||||
## Сводка — что нарушено и насколько срочно
|
||||
|
||||
| # | Нарушение | Срочность | Зона воздействия |
|
||||
|---|---|---|---|
|
||||
| 1 | `@gin.configurable + @dataclass` на `TrainConfigGTAUAV` | 🔴 критично | блокирует всё остальное |
|
||||
| 2 | `@gin.configurable` на `InfoNCELoss`, `WeightedInfoNCELoss`, `GeoLocCaptionDataset` | 🔴 критично | активный источник тихих багов |
|
||||
| 3 | `argparse` в тренировочном коде | 🔴 критично | три источника правды на параметр |
|
||||
| 4 | `include` между `.gin` файлами | 🟡 пограничный | формально нарушает «нет ссылок между конфигами» |
|
||||
| 5 | Один мега-конфиг (нет разделения на оси) | 🟡 следствие #1 | разрешится с #1 |
|
||||
| 6 | Module-level пути в 4 файлах | 🟡 нарушение, но не критично | дубли — реальная проблема |
|
||||
| 7 | `argparse` в скриптах препроцессинга | ⚪ открытый вопрос | спорный случай |
|
||||
| 8 | `@dataclass` в `SOFIAv1Config` (без gin) | ⚪ открытый вопрос | формально не нарушение |
|
||||
| 9 | `@torch.no_grad()` вместо `@torch.inference_mode()` | 🟢 мелочь | косметика |
|
||||
| 10 | `_atomic_save` без cleanup .tmp | 🟢 мелочь | редкое последствие |
|
||||
|
||||
---
|
||||
|
||||
## Что предлагаю обсудить дальше (Шаг 2)
|
||||
|
||||
Прежде чем двигать код, нужно принять решения по 4 развилкам:
|
||||
|
||||
1. **`include` между `.gin`** — терпим как удобный DRY-механизм или приводим к плоской иерархии (каждый эксперимент = самодостаточный набор `.gin` без `include`)?
|
||||
|
||||
2. **Разделение `TrainConfigGTAUAV` на классы** — на сколько и по каким осям?
|
||||
- Вариант **5 классов** (Pipeline / Hardware / Models / Training / Tracking) из «`Рекомендуемые_gin-config_категории.md`».
|
||||
- Вариант **6 классов** (отдельный Loss + Sampler).
|
||||
- Вариант **3 класса** (всё, что было — в Training, плюс Pipeline + Tracking).
|
||||
|
||||
3. **Sofia + dataclass** — оставлять `SOFIAv1Config` как dataclass-структуру внутри модели или переписать в обычный класс/влить в `ModelsConfig`?
|
||||
|
||||
4. **Скрипты** — переводить на gin или оставить argparse?
|
||||
|
||||
После решения этих 4 пунктов план рефакторинга становится однозначным.
|
||||
@@ -25,7 +25,7 @@ from tqdm import tqdm
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.filter_seg")
|
||||
|
||||
SEGM_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-aug/segm")
|
||||
SEGM_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm")
|
||||
EXCLUDE_CLASSES = {0, 4} # background, water
|
||||
DEFAULT_THRESHOLD = 0.90
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ import coloredlogs
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.make_split")
|
||||
|
||||
_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
|
||||
_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
||||
11
scripts/test_dcn.py
Normal file
11
scripts/test_dcn.py
Normal file
@@ -0,0 +1,11 @@
|
||||
import torch
|
||||
from DCNv4 import DCNv4
|
||||
|
||||
dcn = DCNv4(channels=64, group=4).cuda().eval()
|
||||
print('--- 50 forward calls, no_grad ---')
|
||||
with torch.no_grad():
|
||||
for s in range(50):
|
||||
x = torch.randn(8, 4096, 64, device='cuda')
|
||||
_ = dcn(x)
|
||||
if s % 10 == 0:
|
||||
print(f'{s}: {torch.cuda.memory_allocated() / 1e6:.1f} MB')
|
||||
54
src/conf/__init__.py
Normal file
54
src/conf/__init__.py
Normal file
@@ -0,0 +1,54 @@
|
||||
"""Gin-configurable settings for the caption-test project.
|
||||
|
||||
Five universal axes of variability:
|
||||
- PipelineConfig — paths, training schedule, output, resume
|
||||
- HardwareConfig — batch size, accumulation, AMP, gradient checkpointing
|
||||
- TrainingConfig — loss + optimizer + sampler (the recipe)
|
||||
- TrackingConfig — wandb / tensorboard / gradcam / profiler
|
||||
|
||||
Plus a model-family config: ModelsCommonConfig describes the active backbone,
|
||||
and one of {DINOv3, StripNet, SOFIAv1, SOFIAv71} ModelsConfig classes
|
||||
parameterises it.
|
||||
|
||||
Plus PreprocessConfig used only by scripts/make_split.py and
|
||||
scripts/filter_segmentation.py.
|
||||
|
||||
All configs are loaded together via load_all_configs(path2cfg) — see
|
||||
config_loader.py.
|
||||
"""
|
||||
|
||||
from src.conf.config_loader import load_all_configs
|
||||
from src.conf.hardware_conf import HardwareConfig, get_hardware_cfg
|
||||
from src.conf.models_common_conf import ModelsCommonConfig, get_models_common_cfg
|
||||
from src.conf.models_dinov3_conf import DINOv3ModelsConfig, get_models_dinov3_cfg
|
||||
from src.conf.models_sofia_v1_conf import SOFIAv1ModelsConfig, get_models_sofia_v1_cfg
|
||||
from src.conf.models_sofia_v71_conf import SOFIAv71ModelsConfig, get_models_sofia_v71_cfg
|
||||
from src.conf.models_stripnet_conf import StripNetModelsConfig, get_models_stripnet_cfg
|
||||
from src.conf.pipeline_conf import PipelineConfig, get_pipeline_cfg
|
||||
from src.conf.preprocess_conf import PreprocessConfig, get_preprocess_cfg
|
||||
from src.conf.tracking_conf import TrackingConfig, get_tracking_cfg
|
||||
from src.conf.training_conf import TrainingConfig, get_training_cfg
|
||||
|
||||
__all__ = [
|
||||
"DINOv3ModelsConfig",
|
||||
"HardwareConfig",
|
||||
"ModelsCommonConfig",
|
||||
"PipelineConfig",
|
||||
"PreprocessConfig",
|
||||
"SOFIAv1ModelsConfig",
|
||||
"SOFIAv71ModelsConfig",
|
||||
"StripNetModelsConfig",
|
||||
"TrackingConfig",
|
||||
"TrainingConfig",
|
||||
"get_hardware_cfg",
|
||||
"get_models_common_cfg",
|
||||
"get_models_dinov3_cfg",
|
||||
"get_models_sofia_v1_cfg",
|
||||
"get_models_sofia_v71_cfg",
|
||||
"get_models_stripnet_cfg",
|
||||
"get_pipeline_cfg",
|
||||
"get_preprocess_cfg",
|
||||
"get_tracking_cfg",
|
||||
"get_training_cfg",
|
||||
"load_all_configs",
|
||||
]
|
||||
171
src/conf/config_loader.py
Normal file
171
src/conf/config_loader.py
Normal file
@@ -0,0 +1,171 @@
|
||||
"""Single entry point for loading all configs in a training run.
|
||||
|
||||
Layout (per REQUIREMENTS_GIN_STYLE.md §1, extended for multi-experiment):
|
||||
|
||||
in/config_files/
|
||||
├── training.gin # common: shared by all training presets
|
||||
├── hardware_default.gin # common: for DINOv3 / StripNet
|
||||
├── hardware_no_gradckpt.gin # common: for SOFIA backbones
|
||||
├── tracking.gin # common: shared by all
|
||||
└── <preset_name>/
|
||||
├── pipeline.gin # local: output_dir, paths, schedule
|
||||
└── models.gin # local: backbone-specific bindings
|
||||
|
||||
`load_all_configs` parses common files at the path2cfg root + preset locals
|
||||
in `path2cfg/<preset_name>/`. Two-pass loading:
|
||||
Pass 1: read <preset_name>/models.gin to learn the backbone.
|
||||
Pass 2: parse common files + preset locals in one batch.
|
||||
|
||||
Override semantics: if the preset directory contains its own training.gin /
|
||||
hardware.gin / tracking.gin, those files are appended AFTER the common
|
||||
versions, so their bindings win (gin: last write wins).
|
||||
"""
|
||||
|
||||
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_common_conf import ModelsCommonConfig
|
||||
from src.conf.models_dinov3_conf import DINOv3ModelsConfig
|
||||
from src.conf.models_sofia_v1_conf import SOFIAv1ModelsConfig
|
||||
from src.conf.models_sofia_v71_conf import SOFIAv71ModelsConfig
|
||||
from src.conf.models_stripnet_conf import StripNetModelsConfig
|
||||
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__)
|
||||
|
||||
# Maps ModelsCommonConfig.backbone → family-specific config class.
|
||||
_BACKBONE_TO_MODELS_CLS = {
|
||||
"dinov3": DINOv3ModelsConfig,
|
||||
"stripnet": StripNetModelsConfig,
|
||||
"sofia_v1": SOFIAv1ModelsConfig,
|
||||
"sofia_v71": SOFIAv71ModelsConfig,
|
||||
}
|
||||
|
||||
# Sofia backbones disable gradient checkpointing.
|
||||
_NO_GRADCKPT_BACKBONES = {"sofia_v1", "sofia_v71"}
|
||||
|
||||
# Common filenames at the path2cfg root.
|
||||
_COMMON_TRAINING = "training.gin"
|
||||
_COMMON_TRACKING = "tracking.gin"
|
||||
_COMMON_HARDWARE_DEFAULT = "hardware_default.gin"
|
||||
_COMMON_HARDWARE_NO_GRADCKPT = "hardware_no_gradckpt.gin"
|
||||
|
||||
# Files a preset must always have locally (in <path2cfg>/<preset_name>/).
|
||||
_LOCAL_REQUIRED = ("pipeline.gin", "models.gin")
|
||||
|
||||
# Files that can optionally be overridden locally; if present in the preset
|
||||
# directory, they win over the common version.
|
||||
_LOCAL_OVERRIDABLE = ("training.gin", "hardware.gin", "tracking.gin")
|
||||
|
||||
|
||||
def load_all_configs(path2cfg: str, preset_name: str) -> dict[str, Any]:
|
||||
"""Parse common gin files + preset gin files and return all configs.
|
||||
|
||||
Args:
|
||||
path2cfg: Path to in/config_files/ (per REQUIREMENTS_GIN_STYLE.md §5).
|
||||
Must contain training.gin, hardware_default.gin,
|
||||
hardware_no_gradckpt.gin, tracking.gin at its root, plus
|
||||
per-preset subdirectories.
|
||||
preset_name: Name of the preset subdirectory under path2cfg, e.g.
|
||||
'gtauav_balanced'.
|
||||
|
||||
Returns:
|
||||
Dict with keys: 'pipeline', 'hardware', 'training', 'tracking',
|
||||
'models_common', 'models'.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If path2cfg, preset directory, or a required file is missing.
|
||||
ValueError: If models_common.backbone is not a known value.
|
||||
"""
|
||||
common_dir = Path(path2cfg)
|
||||
preset_dir = common_dir / preset_name
|
||||
|
||||
# Sanity checks.
|
||||
if not common_dir.is_dir():
|
||||
raise FileNotFoundError(
|
||||
f"Config root not found: {common_dir}. "
|
||||
f"Per REQUIREMENTS_GIN_STYLE.md §5, this should be "
|
||||
f"<proj_dir>/in/config_files/.",
|
||||
)
|
||||
if not preset_dir.is_dir():
|
||||
raise FileNotFoundError(
|
||||
f"Preset directory not found: {preset_dir}. "
|
||||
f"Available presets in {common_dir}: "
|
||||
f"{sorted(d.name for d in common_dir.iterdir() if d.is_dir())}",
|
||||
)
|
||||
for required in _LOCAL_REQUIRED:
|
||||
if not (preset_dir / required).is_file():
|
||||
raise FileNotFoundError(
|
||||
f"Preset {preset_name} is missing required file '{required}' "
|
||||
f"(looked in {preset_dir})",
|
||||
)
|
||||
|
||||
# ===== Pass 1: peek at models.gin to learn the backbone. =====
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(str(preset_dir / "models.gin"))
|
||||
backbone = ModelsCommonConfig().backbone
|
||||
if backbone not in _BACKBONE_TO_MODELS_CLS:
|
||||
raise ValueError(
|
||||
f"Unknown backbone={backbone!r} in {preset_dir / 'models.gin'}; "
|
||||
f"expected one of {sorted(_BACKBONE_TO_MODELS_CLS)}",
|
||||
)
|
||||
|
||||
# ===== Pass 2: build the full file list and parse in one batch. =====
|
||||
# Order: common first, preset locals last (so locals win on overrides).
|
||||
gin_files: list[Path] = [
|
||||
common_dir / _COMMON_TRAINING,
|
||||
common_dir / (
|
||||
_COMMON_HARDWARE_NO_GRADCKPT if backbone in _NO_GRADCKPT_BACKBONES
|
||||
else _COMMON_HARDWARE_DEFAULT
|
||||
),
|
||||
common_dir / _COMMON_TRACKING,
|
||||
# Always-required preset locals.
|
||||
preset_dir / "pipeline.gin",
|
||||
preset_dir / "models.gin",
|
||||
]
|
||||
|
||||
# Optional preset overrides (rare).
|
||||
for overridable in _LOCAL_OVERRIDABLE:
|
||||
local = preset_dir / overridable
|
||||
if local.is_file():
|
||||
gin_files.append(local)
|
||||
logger.info("Preset %s overrides %s locally", preset_name, overridable)
|
||||
|
||||
# Sanity: all chosen files must exist.
|
||||
for f in gin_files:
|
||||
if not f.is_file():
|
||||
raise FileNotFoundError(f"Required gin file not found: {f}")
|
||||
|
||||
# MANDATORY: clear gin global state before parsing.
|
||||
gin.clear_config()
|
||||
gin.parse_config_files_and_bindings(
|
||||
config_files=[str(f) for f in gin_files],
|
||||
bindings=[],
|
||||
)
|
||||
logger.info(
|
||||
"Loaded preset %s with %d gin files (backbone=%s)",
|
||||
preset_name, len(gin_files), backbone,
|
||||
)
|
||||
|
||||
# Build all configs from gin global state.
|
||||
models_common = ModelsCommonConfig()
|
||||
models_specific = _BACKBONE_TO_MODELS_CLS[models_common.backbone]()
|
||||
|
||||
return {
|
||||
"pipeline": PipelineConfig(),
|
||||
"hardware": HardwareConfig(),
|
||||
"training": TrainingConfig(),
|
||||
"tracking": TrackingConfig(),
|
||||
"models_common": models_common,
|
||||
"models": models_specific,
|
||||
}
|
||||
|
||||
|
||||
46
src/conf/hardware_conf.py
Normal file
46
src/conf/hardware_conf.py
Normal file
@@ -0,0 +1,46 @@
|
||||
"""GPU profile + memory/compute optimisation flags."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class HardwareConfig:
|
||||
"""Hardware-bound parameters: VRAM footprint and throughput.
|
||||
|
||||
These do not change the training recipe (loss/optimizer/sampler), only
|
||||
how many samples fit on the device.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device: str = "cuda",
|
||||
batch_size: int = 8,
|
||||
grad_accum_steps: int = 8,
|
||||
num_workers: int = 4,
|
||||
use_amp: bool = True,
|
||||
gradient_checkpointing: bool = True,
|
||||
reserve_gb: float = 2.0,
|
||||
total_vram_gb: float = 24.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
|
||||
self.total_vram_gb = total_vram_gb
|
||||
# Derived.
|
||||
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 — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}hardware.gin")
|
||||
return HardwareConfig()
|
||||
|
||||
|
||||
50
src/conf/models_common_conf.py
Normal file
50
src/conf/models_common_conf.py
Normal file
@@ -0,0 +1,50 @@
|
||||
"""Backbone-agnostic model parameters.
|
||||
|
||||
`backbone` selects which family-specific Models config is loaded by
|
||||
config_loader.load_all_configs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class ModelsCommonConfig:
|
||||
"""Shared model fields across all backbones.
|
||||
|
||||
`backbone` is the dispatch key — one of:
|
||||
- 'dinov3' → DINOv3ModelsConfig
|
||||
- 'stripnet' → StripNetModelsConfig
|
||||
- 'sofia_v1' → SOFIAv1ModelsConfig
|
||||
- 'sofia_v71' → SOFIAv71ModelsConfig
|
||||
|
||||
`baseline_mode=True` disables text fusion entirely (gates locked at 1.0,
|
||||
DGTRS-CLIP not loaded, TextFusionMLP not built). Used for Δ R@1 baselines.
|
||||
|
||||
`init_gate` controls the initial sigmoid value of GatedFusion gates
|
||||
(0.7 = 70% image, 30% text by default; 0.3 = text-heavy; 0.9 = image-heavy).
|
||||
|
||||
`lrsclip_path` is the path to the DGTRS-CLIP checkpoint (only loaded when
|
||||
text fusion is active).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
backbone: str = "dinov3",
|
||||
baseline_mode: bool = False,
|
||||
init_gate: float = 0.7,
|
||||
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
|
||||
) -> None:
|
||||
self.backbone = backbone
|
||||
self.baseline_mode = baseline_mode
|
||||
self.init_gate = init_gate
|
||||
self.lrsclip_path = lrsclip_path
|
||||
|
||||
|
||||
def get_models_common_cfg(path2cfg: str) -> ModelsCommonConfig:
|
||||
"""Load ONLY models_common config (TESTING ONLY — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}models.gin")
|
||||
return ModelsCommonConfig()
|
||||
|
||||
44
src/conf/models_dinov3_conf.py
Normal file
44
src/conf/models_dinov3_conf.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""DINOv3 backbone configuration: encoders + MONA adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class DINOv3ModelsConfig:
|
||||
"""DINOv3 ViT-L/16 with MONA adapters (CVPR 2025).
|
||||
|
||||
`shared_encoder=True` uses a single DINOv3 WEB encoder for both drone and
|
||||
satellite branches (default; ~432M params total). When False, separate WEB
|
||||
(drone) + SAT (satellite) encoders are built (~733M params total, +4-5GB
|
||||
VRAM).
|
||||
|
||||
MONA adapters are injected in the LAST `mona_last_n_blocks` of the 24
|
||||
ViT blocks (default: 12 = top half). Set to 24 for full-capacity asymmetric
|
||||
setup.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
|
||||
dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
|
||||
shared_encoder: bool = True,
|
||||
mona_bottleneck: int = 64,
|
||||
mona_last_n_blocks: int = 12,
|
||||
lora_rank: int = 4,
|
||||
) -> None:
|
||||
self.dino_web_path = dino_web_path
|
||||
self.dino_sat_path = dino_sat_path
|
||||
self.shared_encoder = shared_encoder
|
||||
self.mona_bottleneck = mona_bottleneck
|
||||
self.mona_last_n_blocks = mona_last_n_blocks
|
||||
self.lora_rank = lora_rank
|
||||
|
||||
|
||||
def get_models_dinov3_cfg(path2cfg: str) -> DINOv3ModelsConfig:
|
||||
"""Load ONLY DINOv3 models config (TESTING ONLY — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}models.gin")
|
||||
return DINOv3ModelsConfig()
|
||||
|
||||
69
src/conf/models_sofia_v1_conf.py
Normal file
69
src/conf/models_sofia_v1_conf.py
Normal file
@@ -0,0 +1,69 @@
|
||||
"""SOFIA v1 backbone configuration: 4-stage StripNet+DCNv4 + GGeM heads."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class SOFIAv1ModelsConfig:
|
||||
"""SOFIA v1: lightweight StripNet+DCNv4 backbone + heads.
|
||||
|
||||
`variant_label` chooses backbone size (architecture dimensions are
|
||||
resolved inside the model code from this label):
|
||||
tiny_tiny: dims [16, 32, 80, 128] (~0.4M)
|
||||
tiny : dims [32, 64, 128, 256] (~1M)
|
||||
small : dims [64, 128, 320, 512] (~3M, default in code)
|
||||
small_v2 : dims [64, 128, 256, 384] (~2M)
|
||||
|
||||
`dcn_variant`: 'v2' = torchvision DeformConv2d (stable). 'v4' = OpenGVLab
|
||||
DCNv4 (faster but ~9 MB / forward leak from C++ extension).
|
||||
|
||||
Text fusion is two-level:
|
||||
- Mid-level: Text-FiLM modulates feature maps before GGeM (when
|
||||
use_text_film_uav / use_text_film_sat = True).
|
||||
- Late-level: GatedFusion on descriptors (handled outside this config).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# ---- Backbone ----
|
||||
variant_label: str = "small", # 'tiny_tiny' | 'tiny' | 'small' | 'small_v2'
|
||||
in_channels: int = 3,
|
||||
input_size: int = 256,
|
||||
dcn_variant: str = "v2", # 'v2' | 'v4'
|
||||
# ---- Heads ----
|
||||
d_descriptor: int = 1024,
|
||||
return_normalized: bool = False,
|
||||
# ---- Altitude-FiLM (UAV head) ----
|
||||
use_film_altitude: bool = True,
|
||||
altitude_norm: float = 500.0,
|
||||
# ---- Text-FiLM (mid-level fusion) ----
|
||||
use_text_film_uav: bool = True,
|
||||
use_text_film_sat: bool = True,
|
||||
text_film_dim: int = 1024,
|
||||
text_film_hidden: int = 256,
|
||||
# ---- LoRA on DGTRS-CLIP text encoder ----
|
||||
lora_rank: int = 4,
|
||||
) -> None:
|
||||
self.variant_label = variant_label
|
||||
self.in_channels = in_channels
|
||||
self.input_size = input_size
|
||||
self.dcn_variant = dcn_variant
|
||||
self.d_descriptor = d_descriptor
|
||||
self.return_normalized = return_normalized
|
||||
self.use_film_altitude = use_film_altitude
|
||||
self.altitude_norm = altitude_norm
|
||||
self.use_text_film_uav = use_text_film_uav
|
||||
self.use_text_film_sat = use_text_film_sat
|
||||
self.text_film_dim = text_film_dim
|
||||
self.text_film_hidden = text_film_hidden
|
||||
self.lora_rank = lora_rank
|
||||
|
||||
|
||||
def get_models_sofia_v1_cfg(path2cfg: str) -> SOFIAv1ModelsConfig:
|
||||
"""Load ONLY SOFIA v1 models config (TESTING ONLY — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}models.gin")
|
||||
return SOFIAv1ModelsConfig()
|
||||
|
||||
178
src/conf/models_sofia_v71_conf.py
Normal file
178
src/conf/models_sofia_v71_conf.py
Normal file
@@ -0,0 +1,178 @@
|
||||
"""SOFIA v7.1 backbone: 4-stage StripDCN + MambaVision + CVGL-Aware Head."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class SOFIAv71ModelsConfig:
|
||||
"""SOFIA v7.1 student model.
|
||||
|
||||
Mirrors src/models/sofia_v71/config.py::SOFIAConfig with one
|
||||
difference: `mamba_extra_kwargs` (a dict in the dataclass) is flattened
|
||||
into 5 explicit fields here, and reassembled into a dict for downstream
|
||||
code.
|
||||
|
||||
Variant scale presets (see model code):
|
||||
Tiny: stem=16/32, dims=[48, 96, 176, 224], depths=[2, 3, 4, 2] (~5M)
|
||||
M : stem=64/128, dims=[256, 512, 1280, 1536], depths=[3, 4, 15, 3] (~500M, default)
|
||||
L : stem=64/128, dims=[256, 512, 1536, 2048], depths=[3, 4, 20, 3] (~1B)
|
||||
|
||||
For the active experiment (Tiny preset, see `presets/gtauav_balanced_sofia/`)
|
||||
you can override individual fields directly without resorting to a
|
||||
'preset' string parameter — every architectural dimension is bindable.
|
||||
|
||||
Tiny needs `num_heads_*=4` (channels 176/224 not divisible by 8) and
|
||||
`mamba_headdim=16` (channels not divisible by 64).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# ---- Variant label (informational, used in logs/checkpoints) ----
|
||||
variant_label: str = "M", # 'M' | 'L' | 'Tiny'
|
||||
# ---- Input ----
|
||||
input_size: int = 256,
|
||||
in_channels: int = 3,
|
||||
# ---- Stem ----
|
||||
stem_mid: int = 64,
|
||||
stem_out: int = 128,
|
||||
# ---- Backbone dimensions (per stage s1..s4) ----
|
||||
# Lists default to None; concrete defaults are filled in __init__ to
|
||||
# avoid the def f(x=[]) anti-pattern.
|
||||
embed_dims: list[int] | None = None, # default [256, 512, 1280, 1536] (M)
|
||||
depths: list[int] | None = None, # default [3, 4, 15, 3] (M)
|
||||
# ---- Stage 1-2 block params ----
|
||||
mbconv_expand: int = 4,
|
||||
se_ratio: int = 16,
|
||||
strip_kernel_s1: int = 7,
|
||||
strip_kernel_s2: int = 5,
|
||||
mix_kernels: list[int] | None = None, # default [3, 5, 7]
|
||||
use_dcn_strip: bool = True,
|
||||
# ---- Stage 3-4 (MambaVision) ----
|
||||
mamba_d_state: int = 16,
|
||||
mamba_dt_rank: int | None = None, # auto = max(1, C // 16)
|
||||
mamba_backend: str = "auto", # 'auto' | 'torch' | 'mamba_ssm'
|
||||
mamba_variant: str = "mamba2", # 'mamba1' | 'mamba2' | 'efficient_vmamba'
|
||||
# mamba_extra_kwargs flattened (assembled back into a dict in __init__):
|
||||
mamba_d_state_mamba2: int = 64,
|
||||
mamba_headdim: int = 64,
|
||||
mamba_expand: int = 2,
|
||||
mamba_d_conv: int = 4,
|
||||
mamba_n_directions: int = 2,
|
||||
# ---- Heads / attention ----
|
||||
num_heads_s3: int = 8,
|
||||
num_heads_s4: int = 8,
|
||||
use_strip_branch_s3: bool = True,
|
||||
use_strip_branch_s4: bool = False,
|
||||
ffn_expand: int = 4,
|
||||
# ---- EVSS bridge ----
|
||||
use_evss_bridge: bool = False,
|
||||
evss_bridge_locations: list[str] | None = None, # default ['pre_stage3']
|
||||
# ---- Neck ----
|
||||
neck_channels: int = 192,
|
||||
# ---- CVGL-Aware Head v7.1-α ----
|
||||
d_descriptor: int = 512,
|
||||
use_asymmetric_heads: bool = True,
|
||||
chp_rings: int = 8,
|
||||
chp_angles: int = 16,
|
||||
chp_harmonics: int = 4,
|
||||
use_film_altitude: bool = True,
|
||||
altitude_norm: float = 500.0,
|
||||
ring_count: int = 4,
|
||||
use_ring_aux: bool = True,
|
||||
# ---- Text fusion ----
|
||||
return_normalized: bool = True,
|
||||
use_text_film_sat: bool = False,
|
||||
use_text_film_uav: bool = False,
|
||||
text_film_dim: int = 1024,
|
||||
text_film_hidden: int = 256,
|
||||
# ---- Weight-sharing ----
|
||||
share_stages_1_2: bool = True,
|
||||
# ---- KD taps ----
|
||||
enable_kd_taps: bool = True,
|
||||
# ---- Deployment ----
|
||||
precision: str = "fp16", # 'fp32' | 'fp16' | 'int8_mixed'
|
||||
# ---- LoRA on DGTRS-CLIP text encoder ----
|
||||
lora_rank: int = 4,
|
||||
) -> None:
|
||||
# Variant label.
|
||||
self.variant_label = variant_label
|
||||
# Input.
|
||||
self.input_size = input_size
|
||||
self.in_channels = in_channels
|
||||
# Stem.
|
||||
self.stem_mid = stem_mid
|
||||
self.stem_out = stem_out
|
||||
# Backbone dimensions.
|
||||
self.embed_dims = embed_dims if embed_dims is not None else [256, 512, 1280, 1536]
|
||||
self.depths = depths if depths is not None else [3, 4, 15, 3]
|
||||
# Stage 1-2.
|
||||
self.mbconv_expand = mbconv_expand
|
||||
self.se_ratio = se_ratio
|
||||
self.strip_kernel_s1 = strip_kernel_s1
|
||||
self.strip_kernel_s2 = strip_kernel_s2
|
||||
self.mix_kernels = mix_kernels if mix_kernels is not None else [3, 5, 7]
|
||||
self.use_dcn_strip = use_dcn_strip
|
||||
# Stage 3-4.
|
||||
self.mamba_d_state = mamba_d_state
|
||||
self.mamba_dt_rank = mamba_dt_rank
|
||||
self.mamba_backend = mamba_backend
|
||||
self.mamba_variant = mamba_variant
|
||||
self.mamba_d_state_mamba2 = mamba_d_state_mamba2
|
||||
self.mamba_headdim = mamba_headdim
|
||||
self.mamba_expand = mamba_expand
|
||||
self.mamba_d_conv = mamba_d_conv
|
||||
self.mamba_n_directions = mamba_n_directions
|
||||
# Heads.
|
||||
self.num_heads_s3 = num_heads_s3
|
||||
self.num_heads_s4 = num_heads_s4
|
||||
self.use_strip_branch_s3 = use_strip_branch_s3
|
||||
self.use_strip_branch_s4 = use_strip_branch_s4
|
||||
self.ffn_expand = ffn_expand
|
||||
# EVSS.
|
||||
self.use_evss_bridge = use_evss_bridge
|
||||
self.evss_bridge_locations = (
|
||||
evss_bridge_locations if evss_bridge_locations is not None else ["pre_stage3"]
|
||||
)
|
||||
# Neck.
|
||||
self.neck_channels = neck_channels
|
||||
# CVGL Head.
|
||||
self.d_descriptor = d_descriptor
|
||||
self.use_asymmetric_heads = use_asymmetric_heads
|
||||
self.chp_rings = chp_rings
|
||||
self.chp_angles = chp_angles
|
||||
self.chp_harmonics = chp_harmonics
|
||||
self.use_film_altitude = use_film_altitude
|
||||
self.altitude_norm = altitude_norm
|
||||
self.ring_count = ring_count
|
||||
self.use_ring_aux = use_ring_aux
|
||||
# Text fusion.
|
||||
self.return_normalized = return_normalized
|
||||
self.use_text_film_sat = use_text_film_sat
|
||||
self.use_text_film_uav = use_text_film_uav
|
||||
self.text_film_dim = text_film_dim
|
||||
self.text_film_hidden = text_film_hidden
|
||||
# Sharing / KD / deploy.
|
||||
self.share_stages_1_2 = share_stages_1_2
|
||||
self.enable_kd_taps = enable_kd_taps
|
||||
self.precision = precision
|
||||
# LoRA.
|
||||
self.lora_rank = lora_rank
|
||||
# Derived: assemble mamba_extra_kwargs back for downstream consumers.
|
||||
self.mamba_extra_kwargs = {
|
||||
"d_state_mamba2": self.mamba_d_state_mamba2,
|
||||
"headdim": self.mamba_headdim,
|
||||
"expand": self.mamba_expand,
|
||||
"d_conv": self.mamba_d_conv,
|
||||
"n_directions": self.mamba_n_directions,
|
||||
}
|
||||
|
||||
|
||||
def get_models_sofia_v71_cfg(path2cfg: str) -> SOFIAv71ModelsConfig:
|
||||
"""Load ONLY SOFIA v71 models config (TESTING ONLY — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}models.gin")
|
||||
return SOFIAv71ModelsConfig()
|
||||
|
||||
|
||||
40
src/conf/models_stripnet_conf.py
Normal file
40
src/conf/models_stripnet_conf.py
Normal file
@@ -0,0 +1,40 @@
|
||||
"""StripNet backbone configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class StripNetModelsConfig:
|
||||
"""StripNet-small encoder with Conv-MONA adaptation.
|
||||
|
||||
`stripnet_freeze=True` keeps the backbone frozen and only trains MONA on
|
||||
the last `stripnet_mona_last_n_stages` of 4 stages.
|
||||
|
||||
`stripnet_freeze=False` (full fine-tune) makes the backbone trainable; in
|
||||
that case backbone params get a separate LR group at
|
||||
`learning_rate * stripnet_backbone_lr_factor` (typically 0.1).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
|
||||
stripnet_freeze: bool = True,
|
||||
stripnet_mona_last_n_stages: int = 2,
|
||||
stripnet_backbone_lr_factor: float = 0.1,
|
||||
lora_rank: int = 4,
|
||||
) -> None:
|
||||
self.stripnet_path = stripnet_path
|
||||
self.stripnet_freeze = stripnet_freeze
|
||||
self.stripnet_mona_last_n_stages = stripnet_mona_last_n_stages
|
||||
self.stripnet_backbone_lr_factor = stripnet_backbone_lr_factor
|
||||
self.lora_rank = lora_rank
|
||||
|
||||
|
||||
def get_models_stripnet_cfg(path2cfg: str) -> StripNetModelsConfig:
|
||||
"""Load ONLY StripNet models config (TESTING ONLY — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}models.gin")
|
||||
return StripNetModelsConfig()
|
||||
|
||||
53
src/conf/pipeline_conf.py
Normal file
53
src/conf/pipeline_conf.py
Normal file
@@ -0,0 +1,53 @@
|
||||
|
||||
"""Pipeline orchestration: data IO, training schedule, output, resume."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class PipelineConfig:
|
||||
"""What to train on, where to save, and how long.
|
||||
|
||||
All paths are absolute or relative to the project root. Defaults match
|
||||
the servml workstation; override in pipeline.gin for other machines.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# Data inputs.
|
||||
train_json: str = "meta/train_80.json",
|
||||
test_json: str = "meta/test_20.json",
|
||||
rgb_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR",
|
||||
caption_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_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 — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}pipeline.gin")
|
||||
return PipelineConfig()
|
||||
|
||||
|
||||
56
src/conf/preprocess_conf.py
Normal file
56
src/conf/preprocess_conf.py
Normal file
@@ -0,0 +1,56 @@
|
||||
"""Preprocessing configuration: train/test split + segmentation filter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class PreprocessConfig:
|
||||
"""Used only by scripts/make_split.py and scripts/filter_segmentation.py.
|
||||
|
||||
Lives in a separate preset (presets/preprocess/preprocess.gin) — it is
|
||||
not consumed by the training pipeline. Held independently from
|
||||
PipelineConfig so that preprocess can run without a training preset
|
||||
being active.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# Inputs.
|
||||
rgb_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR",
|
||||
segm_root: str = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-aug/segm",
|
||||
# make_split.py params.
|
||||
split_ratio: float = 0.8,
|
||||
split_seed: int = 42,
|
||||
split_input_train: str = "cross-area-drone2sate-train.json",
|
||||
split_input_test: str = "cross-area-drone2sate-test.json",
|
||||
split_output_dir: str = "meta",
|
||||
split_output_train: str = "train_80.json",
|
||||
split_output_test: str = "test_20.json",
|
||||
# filter_segmentation.py params.
|
||||
seg_threshold: float = 0.90,
|
||||
seg_exclude_classes: list[int] | None = None, # default [0, 4]: background + water
|
||||
seg_filter_output: str = "meta/seg_filter.json",
|
||||
) -> None:
|
||||
self.rgb_root = rgb_root
|
||||
self.segm_root = segm_root
|
||||
self.split_ratio = split_ratio
|
||||
self.split_seed = split_seed
|
||||
self.split_input_train = split_input_train
|
||||
self.split_input_test = split_input_test
|
||||
self.split_output_dir = split_output_dir
|
||||
self.split_output_train = split_output_train
|
||||
self.split_output_test = split_output_test
|
||||
self.seg_threshold = seg_threshold
|
||||
self.seg_exclude_classes = seg_exclude_classes if seg_exclude_classes is not None else [0, 4]
|
||||
self.seg_filter_output = seg_filter_output
|
||||
|
||||
|
||||
def get_preprocess_cfg(path2cfg: str) -> PreprocessConfig:
|
||||
"""Load preprocess config from the preprocess preset directory."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}preprocess.gin")
|
||||
return PreprocessConfig()
|
||||
|
||||
|
||||
51
src/conf/tracking_conf.py
Normal file
51
src/conf/tracking_conf.py
Normal file
@@ -0,0 +1,51 @@
|
||||
"""Experiment tracking + diagnostics.
|
||||
|
||||
Independent axis: changing these flags does not affect training results,
|
||||
only what is observed/recorded.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class TrackingConfig:
|
||||
"""Wandb / TensorBoard / Grad-CAM / profiler / gradient norms."""
|
||||
|
||||
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 — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}tracking.gin")
|
||||
return TrackingConfig()
|
||||
|
||||
|
||||
93
src/conf/training_conf.py
Normal file
93
src/conf/training_conf.py
Normal file
@@ -0,0 +1,93 @@
|
||||
"""Training recipe: loss + optimizer + sampler.
|
||||
|
||||
Three concerns kept together because they form one coherent recipe — they
|
||||
co-vary across experiments. Splitting Loss vs Optimizer vs Sampler can be
|
||||
done later if a need emerges.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gin
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class TrainingConfig:
|
||||
"""Loss + optimizer + sampler.
|
||||
|
||||
Selects between InfoNCELoss and WeightedInfoNCELoss via `loss_type`.
|
||||
Selects between DSS / mutex / plain shuffle via `sampler_type`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# ---- Loss: shared between InfoNCELoss and WeightedInfoNCELoss ----
|
||||
loss_type: str = "symmetric", # 'symmetric' | 'weighted'
|
||||
tau_init: float = 0.07,
|
||||
tau_min: float = 0.01,
|
||||
tau_max: float = 0.1,
|
||||
learnable_temperature: bool = True,
|
||||
label_smoothing: float = 0.1,
|
||||
# ---- Loss: InfoNCELoss-only ----
|
||||
tau_final: float = 0.01, # cosine-schedule final tau (when not learnable)
|
||||
weight_q2g: float = 0.6,
|
||||
weight_g2q: float = 0.4,
|
||||
hard_mining_k: int = 0,
|
||||
neg_bank_size: int = 0,
|
||||
# ---- Loss: WeightedInfoNCELoss-only ----
|
||||
weighted_loss_k: float = 5.0, # sigmoid steepness for weight→eps mapping
|
||||
# ---- Optimizer ----
|
||||
learning_rate: float = 1e-4,
|
||||
text_lr_factor: float = 0.1, # lr * factor for DGTRS-CLIP/LoRA params
|
||||
weight_decay: float = 1e-4,
|
||||
grad_clip: float = 1.0,
|
||||
# ---- Sampler ----
|
||||
sampler_type: str = "mutex", # 'mutex' | 'dss' | 'none'
|
||||
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,
|
||||
# Legacy alias (kept until train_gtauav.py is rewritten in step 4).
|
||||
use_mutex_sampler: bool = True,
|
||||
) -> None:
|
||||
# Loss (shared).
|
||||
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
|
||||
# Loss (InfoNCE-specific).
|
||||
self.tau_final = tau_final
|
||||
self.weight_q2g = weight_q2g
|
||||
self.weight_g2q = weight_g2q
|
||||
self.hard_mining_k = hard_mining_k
|
||||
self.neg_bank_size = neg_bank_size
|
||||
# Loss (WeightedInfoNCE-specific).
|
||||
self.weighted_loss_k = weighted_loss_k
|
||||
# Optimizer.
|
||||
self.learning_rate = learning_rate
|
||||
self.text_lr_factor = text_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
|
||||
self.use_mutex_sampler = use_mutex_sampler
|
||||
|
||||
|
||||
def get_training_cfg(path2cfg: str) -> TrainingConfig:
|
||||
"""Load ONLY training config (TESTING ONLY — production uses load_all_configs)."""
|
||||
gin.clear_config()
|
||||
gin.parse_config_file(f"{path2cfg}training.gin")
|
||||
return TrainingConfig()
|
||||
|
||||
|
||||
@@ -30,13 +30,17 @@ LOGGER = logging.getLogger("caption_test.gtauav_dataset")
|
||||
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
|
||||
# Default paths.
|
||||
_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
|
||||
_CAPTION_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-captions")
|
||||
_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
|
||||
_CAPTION_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions")
|
||||
_EMPTY_CAPTION = ""
|
||||
|
||||
# Regex to split P1/P2/P3 sections.
|
||||
_P_SPLIT = re.compile(r"\*\*P[123][^*]*\*\*\s*:?\s*")
|
||||
|
||||
# TODO: Transforms
|
||||
# sat_transforms = ...
|
||||
# drone_transforms = ...
|
||||
|
||||
|
||||
def _parse_caption_levels(output: str) -> tuple[str, str, str]:
|
||||
"""Split VLM caption output into L1, L2, L3 levels.
|
||||
@@ -104,6 +108,10 @@ class GTAUAVDataset(Dataset):
|
||||
image_transform: Fallback single transform for both (used if drone/sat not set).
|
||||
filter_meta: Path to seg_filter.json (exclude 90%+ bg/water).
|
||||
drop_caption_prob: Probability of dropping captions (ablation).
|
||||
meta_csvs: Optional list of CSV paths with `img_name,drone_height,...`
|
||||
columns (e.g. `cross-area-drone2sate-{train,test}_drone_meta.csv`).
|
||||
When provided, an `altitude` (meters) field is attached per drone
|
||||
sample. Defaults to scanning the rgb_root for *_drone_meta.csv.
|
||||
seed: Random seed.
|
||||
"""
|
||||
|
||||
@@ -117,6 +125,7 @@ class GTAUAVDataset(Dataset):
|
||||
image_transform: Callable[[Image.Image], torch.Tensor] | None = None,
|
||||
filter_meta: str | None = None,
|
||||
drop_caption_prob: float = 0.0,
|
||||
meta_csvs: list[str] | None = None,
|
||||
seed: int = 0,
|
||||
) -> None:
|
||||
self.rgb_root = Path(rgb_root)
|
||||
@@ -131,6 +140,9 @@ class GTAUAVDataset(Dataset):
|
||||
if filter_meta is not None:
|
||||
self._load_filter(Path(filter_meta))
|
||||
|
||||
# Load drone altitude index (img_name -> meters). Empty dict if no CSVs.
|
||||
self.altitude_index: dict[str, float] = self._load_altitude_index(meta_csvs)
|
||||
|
||||
# Load caption index.
|
||||
LOGGER.info("📚 Loading caption index from %s", caption_root)
|
||||
self.caption_index = _load_caption_index(self.caption_root)
|
||||
@@ -141,6 +153,43 @@ class GTAUAVDataset(Dataset):
|
||||
self._load_pairs(Path(pair_json))
|
||||
LOGGER.info("✅ Loaded %d pairs from %s", len(self.entries), pair_json)
|
||||
|
||||
def _load_altitude_index(self, meta_csvs: list[str] | None) -> dict[str, float]:
|
||||
"""Build {drone_img_name: altitude_meters} from drone_meta CSVs.
|
||||
|
||||
If `meta_csvs` is None, auto-discovers `*_drone_meta.csv` (TSV format,
|
||||
columns img_name/drone_height/...) under `self.rgb_root`. Missing or
|
||||
unreadable files are skipped silently — altitude defaults to 0.0
|
||||
downstream when an entry is missing.
|
||||
"""
|
||||
if meta_csvs is None:
|
||||
meta_csvs = [str(p) for p in sorted(self.rgb_root.glob("*_drone_meta.csv"))]
|
||||
# Prefer `*_drone_meta_new.csv` if present (overrides original).
|
||||
new_csvs = [str(p) for p in sorted(self.rgb_root.glob("*_drone_meta_new.csv"))]
|
||||
meta_csvs = new_csvs or meta_csvs
|
||||
index: dict[str, float] = {}
|
||||
for csv_path in meta_csvs:
|
||||
path = Path(csv_path)
|
||||
if not path.exists():
|
||||
continue
|
||||
try:
|
||||
with path.open() as f:
|
||||
header = f.readline().rstrip("\n").split("\t")
|
||||
name_idx = header.index("img_name")
|
||||
height_idx = header.index("drone_height")
|
||||
for line in f:
|
||||
parts = line.rstrip("\n").split("\t")
|
||||
if len(parts) <= max(name_idx, height_idx):
|
||||
continue
|
||||
try:
|
||||
index[parts[name_idx]] = float(parts[height_idx])
|
||||
except ValueError:
|
||||
continue
|
||||
except (OSError, ValueError) as exc:
|
||||
LOGGER.warning("Failed to parse drone meta CSV %s: %s", path, exc)
|
||||
if index:
|
||||
LOGGER.info("📐 Altitude index: %d drones (from %d CSV)", len(index), len(meta_csvs))
|
||||
return index
|
||||
|
||||
def _load_filter(self, path: Path) -> None:
|
||||
with open(path) as f:
|
||||
meta = json.load(f)
|
||||
@@ -200,6 +249,7 @@ class GTAUAVDataset(Dataset):
|
||||
"caption_l2": l2,
|
||||
"caption_l3": l3,
|
||||
"sat_captions": sat_captions,
|
||||
"altitude": self.altitude_index.get(drone_name, 0.0),
|
||||
})
|
||||
|
||||
def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor:
|
||||
@@ -267,6 +317,7 @@ class GTAUAVDataset(Dataset):
|
||||
"pair_id": entry["drone_name"],
|
||||
"sat_name": sat_name,
|
||||
"positive_weight": pos_weight,
|
||||
"altitude": float(entry["altitude"]),
|
||||
}
|
||||
|
||||
|
||||
@@ -286,6 +337,7 @@ def collate_gtauav_batch(
|
||||
"pair_ids": [b["pair_id"] for b in batch],
|
||||
"sat_names": [b["sat_name"] for b in batch],
|
||||
"positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32),
|
||||
"altitude": torch.tensor([b["altitude"] for b in batch], dtype=torch.float32),
|
||||
}
|
||||
|
||||
|
||||
@@ -371,6 +423,7 @@ class GTAUAVDroneQuery(Dataset):
|
||||
"caption_l2": entry["caption_l2"],
|
||||
"caption_l3": entry["caption_l3"],
|
||||
"valid_sat_names": list(entry["sat_candidates"]),
|
||||
"altitude": float(entry.get("altitude", 0.0)),
|
||||
}
|
||||
|
||||
|
||||
@@ -392,4 +445,5 @@ def collate_drone_query(batch: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
"caption_l2": [b["caption_l2"] for b in batch],
|
||||
"caption_l3": [b["caption_l3"] for b in batch],
|
||||
"valid_sat_names": [b["valid_sat_names"] for b in batch],
|
||||
"altitude": torch.tensor([b["altitude"] for b in batch], dtype=torch.float32),
|
||||
}
|
||||
|
||||
245
src/eval/evaluator.py
Normal file
245
src/eval/evaluator.py
Normal file
@@ -0,0 +1,245 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Retrieval evaluation for GTA-UAV-LR cross-view geo-localization.
|
||||
|
||||
Computes R@K and MRR for both q→g (drone→satellite) and g→q (satellite→drone)
|
||||
on the full satellite gallery. Multi-match: a query counts as a hit@K if ANY
|
||||
of its valid satellite matches (sat_candidates) appears in the top-K.
|
||||
|
||||
Body transplanted byte-for-byte from src/training/train_gtauav.py::_evaluate
|
||||
in the main branch. The single difference is the type annotation
|
||||
`model: AsymmetricEncoder` → `model: nn.Module` (relaxed for duck-typing
|
||||
across encoder families); semantically identical to the main-branch version.
|
||||
|
||||
Note: not to be confused with src/eval/evaluate.py (legacy v2 helper for
|
||||
UAV-VisLoc with a different signature). This module lives at
|
||||
src/eval/evaluator.py and is the active evaluator for v3 GTA-UAV-LR.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from src.models.asymmetric_encoder import AsymmetricEncoder
|
||||
from src.datasets.gtauav_dataset import (
|
||||
GTAUAVDataset,
|
||||
GTAUAVDroneQuery,
|
||||
GTAUAVSatGallery,
|
||||
collate_drone_query,
|
||||
collate_sat_gallery,
|
||||
)
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.evaluator")
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(
|
||||
model: AsymmetricEncoder,
|
||||
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.
|
||||
|
||||
Standard CVGL retrieval: forward every unique satellite in the dataset
|
||||
once (gallery), forward every drone query, then rank gallery by
|
||||
cosine similarity. A query counts as a hit@K if ANY of its valid
|
||||
satellite matches (pair_pos_sate_img_list ∪ pair_pos_semipos_sate_img_list)
|
||||
appears in the top-K.
|
||||
|
||||
`max_batches` subsamples the drone queries (not the gallery) — useful
|
||||
for a quick train-side sanity check.
|
||||
|
||||
Args:
|
||||
model: Encoder with `encode_query(drone_img, l1, l2, l3)`
|
||||
and `encode_gallery(sat_img, l1, l2, l3)`. Must expose
|
||||
`fusion_query.gate_value` and `fusion_gallery.gate_value`.
|
||||
loader: DataLoader over a GTAUAVDataset (used only to pull dataset
|
||||
+ batch_size/num_workers/pin_memory; iteration is bypassed —
|
||||
we build separate query and gallery loaders inside).
|
||||
device: Torch device string.
|
||||
loss_fn: If provided, computes per-batch loss against paired gallery
|
||||
entries (uses the first valid sat per query as its positive).
|
||||
The mean loss appears in the returned dict under 'loss'.
|
||||
epoch, total_epochs: Passed through to loss_fn.
|
||||
k_values: K values for R@K (e.g. (1, 5, 10)).
|
||||
max_batches: Cap on query batches for quick sanity checks (gallery
|
||||
is always full).
|
||||
desc: tqdm description prefix.
|
||||
|
||||
Returns:
|
||||
Dict with: r@K_q2g, ap_q2g (= MRR), r@K_g2q, ap_g2q, loss (optional),
|
||||
n_query, n_gallery, n_scored_g2q, gate_q, gate_g.
|
||||
"""
|
||||
dataset = loader.dataset
|
||||
if not isinstance(dataset, GTAUAVDataset):
|
||||
raise TypeError(
|
||||
f"evaluate() expects GTAUAVDataset, got {type(dataset).__name__}",
|
||||
)
|
||||
|
||||
model.eval()
|
||||
|
||||
batch_size = loader.batch_size or 32
|
||||
num_workers = getattr(loader, "num_workers", 0)
|
||||
pin_memory = getattr(loader, "pin_memory", False)
|
||||
|
||||
gallery_ds = GTAUAVSatGallery(dataset)
|
||||
query_ds = GTAUAVDroneQuery(dataset)
|
||||
|
||||
gallery_loader = DataLoader(
|
||||
gallery_ds,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
pin_memory=pin_memory,
|
||||
collate_fn=collate_sat_gallery,
|
||||
)
|
||||
query_loader = DataLoader(
|
||||
query_ds,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
pin_memory=pin_memory,
|
||||
collate_fn=collate_drone_query,
|
||||
)
|
||||
|
||||
# --- Gallery forward (all unique sats) ---
|
||||
gallery_embs: list[torch.Tensor] = []
|
||||
gallery_names: list[str] = []
|
||||
for batch in tqdm(gallery_loader, desc=f" {desc}-gallery", unit="batch", leave=False):
|
||||
sat_img = batch["sat_img"].to(device, non_blocking=True)
|
||||
g = model.encode_gallery(
|
||||
sat_img,
|
||||
batch["sat_caption_l1"], batch["sat_caption_l2"], batch["sat_caption_l3"],
|
||||
)
|
||||
gallery_embs.append(g.cpu())
|
||||
gallery_names.extend(batch["sat_names"])
|
||||
gallery = torch.cat(gallery_embs, dim=0) # [N_sat, D]
|
||||
|
||||
# --- Query forward (optionally subsampled via max_batches) ---
|
||||
query_embs: list[torch.Tensor] = []
|
||||
query_valid_names: list[list[str]] = []
|
||||
batch_losses: list[float] = []
|
||||
sat_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(gallery_names)}
|
||||
|
||||
for i, batch in enumerate(tqdm(query_loader, desc=f" {desc}-query", unit="batch", leave=False)):
|
||||
if max_batches is not None and i >= max_batches:
|
||||
break
|
||||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||||
q = model.encode_query(
|
||||
drone_img,
|
||||
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
|
||||
)
|
||||
query_embs.append(q.cpu())
|
||||
query_valid_names.extend(batch["valid_sat_names"])
|
||||
|
||||
# Per-batch loss: use first valid sat per query as its paired gallery.
|
||||
if loss_fn is not None:
|
||||
pair_indices: list[int] = []
|
||||
for names in batch["valid_sat_names"]:
|
||||
for name in names:
|
||||
if name in sat_name_to_idx:
|
||||
pair_indices.append(sat_name_to_idx[name])
|
||||
break
|
||||
else:
|
||||
pair_indices.append(-1)
|
||||
if all(idx >= 0 for idx in pair_indices):
|
||||
paired_gallery = gallery[pair_indices].to(device)
|
||||
fake_embeddings = {
|
||||
"query": q,
|
||||
"gallery": paired_gallery,
|
||||
"gate_q": model.fusion_query.gate_value,
|
||||
"gate_g": model.fusion_gallery.gate_value,
|
||||
}
|
||||
loss_dict = loss_fn(fake_embeddings, epoch=epoch, total_epochs=total_epochs)
|
||||
batch_losses.append(float(loss_dict["total"].item()))
|
||||
|
||||
query = torch.cat(query_embs, dim=0) # [N_q, D]
|
||||
n_query = query.size(0)
|
||||
|
||||
# --- Similarity + rankings ---
|
||||
sim = query @ gallery.t() # [N_q, N_sat]
|
||||
sorted_idx = sim.argsort(dim=1, descending=True)
|
||||
|
||||
metrics: dict[str, float] = {}
|
||||
if batch_losses:
|
||||
metrics["loss"] = sum(batch_losses) / len(batch_losses)
|
||||
|
||||
# Precompute valid gallery index sets per query.
|
||||
valid_idx_per_query: list[set[int]] = []
|
||||
for names in query_valid_names:
|
||||
valid = {sat_name_to_idx[n] for n in names if n in sat_name_to_idx}
|
||||
valid_idx_per_query.append(valid)
|
||||
|
||||
# R@K with multi-match.
|
||||
for k in k_values:
|
||||
hits = 0
|
||||
for i in range(n_query):
|
||||
top_k = set(sorted_idx[i, :k].tolist())
|
||||
if valid_idx_per_query[i] & top_k:
|
||||
hits += 1
|
||||
metrics[f"r@{k}_q2g"] = hits / max(n_query, 1)
|
||||
|
||||
# MRR over valid matches (kept key `ap_q2g` for CSV/plot compatibility).
|
||||
mrr_sum = 0.0
|
||||
n_scored = 0
|
||||
for i in range(n_query):
|
||||
valid = valid_idx_per_query[i]
|
||||
if not valid:
|
||||
continue
|
||||
n_scored += 1
|
||||
for rank, gidx in enumerate(sorted_idx[i].tolist()):
|
||||
if gidx in valid:
|
||||
mrr_sum += 1.0 / (rank + 1)
|
||||
break
|
||||
metrics["ap_q2g"] = mrr_sum / max(n_scored, 1)
|
||||
|
||||
# --- g2q (satellite → drone): invert ground-truth ---
|
||||
n_gallery = gallery.size(0)
|
||||
valid_q_per_sat: list[set[int]] = [set() for _ in range(n_gallery)]
|
||||
for q_idx, gset in enumerate(valid_idx_per_query):
|
||||
for g_idx in gset:
|
||||
valid_q_per_sat[g_idx].add(q_idx)
|
||||
|
||||
sorted_idx_g2q = sim.t().argsort(dim=1, descending=True) # [N_sat, n_query]
|
||||
n_scored_g2q = sum(1 for s in valid_q_per_sat if s)
|
||||
|
||||
for k in k_values:
|
||||
hits_g2q = 0
|
||||
for i in range(n_gallery):
|
||||
valid = valid_q_per_sat[i]
|
||||
if not valid:
|
||||
continue
|
||||
top_k = set(sorted_idx_g2q[i, :k].tolist())
|
||||
if valid & top_k:
|
||||
hits_g2q += 1
|
||||
metrics[f"r@{k}_g2q"] = hits_g2q / max(n_scored_g2q, 1)
|
||||
|
||||
mrr_sum_g2q = 0.0
|
||||
for i in range(n_gallery):
|
||||
valid = valid_q_per_sat[i]
|
||||
if not valid:
|
||||
continue
|
||||
for rank, qidx in enumerate(sorted_idx_g2q[i].tolist()):
|
||||
if qidx in valid:
|
||||
mrr_sum_g2q += 1.0 / (rank + 1)
|
||||
break
|
||||
metrics["ap_g2q"] = mrr_sum_g2q / max(n_scored_g2q, 1)
|
||||
|
||||
metrics["n_query"] = float(n_query)
|
||||
metrics["n_gallery"] = float(n_gallery)
|
||||
metrics["n_scored_g2q"] = float(n_scored_g2q)
|
||||
|
||||
metrics["gate_q"] = model.fusion_query.gate_value
|
||||
metrics["gate_g"] = model.fusion_gallery.gate_value
|
||||
return metrics
|
||||
|
||||
|
||||
@@ -10,7 +10,6 @@ Supports both learnable temperature (CLIP-style logit_scale) and fixed/scheduled
|
||||
|
||||
import math
|
||||
|
||||
import gin
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
@@ -81,10 +80,13 @@ def cosine_temperature(
|
||||
return tau_final + (tau_init - tau_final) * cosine
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class InfoNCELoss(nn.Module):
|
||||
"""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:
|
||||
temperature_init: Initial temperature value.
|
||||
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 gin
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class WeightedInfoNCELoss(nn.Module):
|
||||
"""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))
|
||||
where w_i is the positive weight (e.g. IoU with matched satellite crop).
|
||||
Higher weight → lower eps → sharper target (strong positive).
|
||||
|
||||
73
src/main.py
Normal file
73
src/main.py
Normal file
@@ -0,0 +1,73 @@
|
||||
"""Entry point: load configs and run training.
|
||||
|
||||
Usage: (commands for each preset)
|
||||
python -m src.main gtauav_balanced
|
||||
python -m src.main gtauav_balanced_asym
|
||||
python -m src.main gtauav_balanced_stripnet
|
||||
python -m src.main gtauav_balanced_stripnet_unfrozen
|
||||
python -m src.main gtauav_baseline
|
||||
python -m src.main gtauav_baseline_asym
|
||||
python -m src.main gtauav_baseline_stripnet
|
||||
python -m src.main gtauav_baseline_stripnet_unfrozen
|
||||
python -m src.main gtauav_image_heavy
|
||||
python -m src.main gtauav_gtauav_text_heavy
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
|
||||
import coloredlogs
|
||||
|
||||
from src.conf.config_loader import load_all_configs
|
||||
from src.training.trainer_new import Trainer
|
||||
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
|
||||
|
||||
# -------------------------------------------------------
|
||||
''' ONLY FOR DEBUG with launch.json config:
|
||||
"args": ["main gtauav_balanced"] -> so need to extract
|
||||
preset name "gtauav_balanced"
|
||||
'''
|
||||
preset_name = preset_name.split(' ')[1]
|
||||
# -------------------------------------------------------
|
||||
|
||||
configs = load_all_configs(path2cfg, preset_name)
|
||||
|
||||
trainer = Trainer(
|
||||
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"],
|
||||
)
|
||||
|
||||
trainer.train()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
@@ -192,23 +192,25 @@ class DINOv3ViT(nn.Module):
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
|
||||
"""Load from .pth or .safetensors checkpoint."""
|
||||
model = cls()
|
||||
path = Path(path)
|
||||
LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
|
||||
if path.suffix == ".safetensors":
|
||||
state = load_safetensors(str(path))
|
||||
else:
|
||||
state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||
if "model" in state:
|
||||
state = state["model"]
|
||||
elif "state_dict" in state:
|
||||
state = state["state_dict"]
|
||||
model.load_state_dict(state, strict=False)
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
|
||||
return model
|
||||
|
||||
try:
|
||||
"""Load from .pth or .safetensors checkpoint."""
|
||||
model = cls()
|
||||
path = Path(path)
|
||||
LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
|
||||
if path.suffix == ".safetensors":
|
||||
state = load_safetensors(str(path))
|
||||
else:
|
||||
state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||
if "model" in state:
|
||||
state = state["model"]
|
||||
elif "state_dict" in state:
|
||||
state = state["state_dict"]
|
||||
model.load_state_dict(state, strict=False)
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
|
||||
return model
|
||||
except FileNotFoundError as e:
|
||||
LOGGER.exception(msg=e.strerror)
|
||||
|
||||
# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
|
||||
# in src/models/dgtrs/model.py (DGTRSTextEncoder)
|
||||
@@ -462,6 +464,7 @@ class AsymmetricEncoder(nn.Module):
|
||||
caption_l1: list[str] | None = None,
|
||||
caption_l2: list[str] | None = None,
|
||||
caption_l3: list[str] | None = None,
|
||||
altitude: torch.Tensor | None = None, # noqa: ARG002 — accepted for API parity with SOFIAFusionEncoder
|
||||
) -> torch.Tensor:
|
||||
"""Encode drone → normalized query embedding with per-sample text mask."""
|
||||
drone_feat = self.encode_drone(drone_img)
|
||||
@@ -492,6 +495,7 @@ class AsymmetricEncoder(nn.Module):
|
||||
sat_caption_l1: list[str] | None = None,
|
||||
sat_caption_l2: list[str] | None = None,
|
||||
sat_caption_l3: list[str] | None = None,
|
||||
altitude: torch.Tensor | None = None, # noqa: ARG002 — accepted for API parity with SOFIAFusionEncoder
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Forward pass.
|
||||
|
||||
|
||||
41
src/models/sofia_v1/__init__.py
Normal file
41
src/models/sofia_v1/__init__.py
Normal file
@@ -0,0 +1,41 @@
|
||||
"""SOFIA v1 — StripNet + DCNv4 hierarchical CNN backbone for CVGL.
|
||||
|
||||
Lightweight 4-stage backbone (~5–30M params depending on variant). Outputs
|
||||
features at 4 scales (last stage is 8x8 for 256x256 input).
|
||||
|
||||
Variants (in `stripnet_model_dcn.VARIANT_MAP`):
|
||||
- `tiny_tiny`: dims [16, 32, 80, 128]
|
||||
- `tiny` : dims [32, 64, 128, 256]
|
||||
- `small` : dims [64, 128, 320, 512] (default)
|
||||
- `small_v2` : dims [64, 128, 256, 384]
|
||||
|
||||
Use `SOFIAv1FusionEncoder` from `src.models.sofia_v1_fusion_encoder` for
|
||||
end-to-end CVGL training with DGTRS-CLIP captions and altitude.
|
||||
"""
|
||||
|
||||
from .config import SOFIAv1Config
|
||||
from .heads import SatHeadV1, UAVHeadV1
|
||||
from .model import SOFIAv1
|
||||
from .stripnet_model_dcn import (
|
||||
StripNetDCN,
|
||||
VARIANT_MAP,
|
||||
build_stripnet_dcn,
|
||||
get_stripnet_dcn_small,
|
||||
get_stripnet_dcn_small_v2,
|
||||
get_stripnet_dcn_tiny,
|
||||
get_stripnet_dcn_tiny_tiny,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"SOFIAv1",
|
||||
"SOFIAv1Config",
|
||||
"SatHeadV1",
|
||||
"UAVHeadV1",
|
||||
"StripNetDCN",
|
||||
"build_stripnet_dcn",
|
||||
"get_stripnet_dcn_small",
|
||||
"get_stripnet_dcn_small_v2",
|
||||
"get_stripnet_dcn_tiny",
|
||||
"get_stripnet_dcn_tiny_tiny",
|
||||
"VARIANT_MAP",
|
||||
]
|
||||
44
src/models/sofia_v1/config.py
Normal file
44
src/models/sofia_v1/config.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""SOFIA v1 configuration.
|
||||
|
||||
Lightweight 4-stage StripNet+DCNv4 backbone variants (`tiny_tiny`/`tiny`/`small`
|
||||
/`small_v2`) plus simple GGeM-based heads with optional altitude-FiLM and
|
||||
text-FiLM modulation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
|
||||
@dataclass
|
||||
class SOFIAv1Config:
|
||||
# -------- Backbone --------
|
||||
variant: Literal["tiny_tiny", "tiny", "small", "small_v2"] = "small"
|
||||
in_channels: int = 3
|
||||
input_size: int = 256
|
||||
# DCN op variant. "v2" (default) uses torchvision DeformConv2d — stable.
|
||||
# "v4" uses OpenGVLab DCNv4 — faster but has known C++-extension memory
|
||||
# leak (~9 MB per forward) that OOMs in long training runs.
|
||||
dcn_variant: Literal["v2", "v4"] = "v2"
|
||||
|
||||
# -------- Heads --------
|
||||
d_descriptor: int = 1024
|
||||
return_normalized: bool = False # False → wrapper handles L2 after gated fusion
|
||||
|
||||
# Altitude-FiLM (UAV head only).
|
||||
use_film_altitude: bool = True
|
||||
altitude_norm: float = 500.0
|
||||
|
||||
# Text-FiLM (mid-level fusion). Both can be toggled independently.
|
||||
use_text_film_uav: bool = True
|
||||
use_text_film_sat: bool = True
|
||||
text_film_dim: int = 1024
|
||||
text_film_hidden: int = 256
|
||||
|
||||
def summary(self) -> str:
|
||||
return (
|
||||
f"SOFIAv1Config(variant={self.variant}, d={self.d_descriptor}, "
|
||||
f"film_alt={self.use_film_altitude}, "
|
||||
f"text_film(sat={self.use_text_film_sat},uav={self.use_text_film_uav}))"
|
||||
)
|
||||
99
src/models/sofia_v1/heads.py
Normal file
99
src/models/sofia_v1/heads.py
Normal file
@@ -0,0 +1,99 @@
|
||||
"""Heads for SOFIA v1 (StripNet+DCNv4 backbone).
|
||||
|
||||
Designed parallel to SOFIA v7.1 heads but lighter — basic GGeM pooling
|
||||
instead of CHP, optional altitude-FiLM (UAV) and text-FiLM (both).
|
||||
|
||||
The heads produce un-normalized D-dim descriptors when `return_normalized=False`,
|
||||
so that a fusion wrapper can blend with text via gated fusion before the final
|
||||
L2 normalization.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from src.models.sofia_v71.layers import AltitudeFiLM, GGeM, TextFiLM
|
||||
|
||||
|
||||
class SatHeadV1(nn.Module):
|
||||
"""Satellite head: [TextFiLM] + GGeM + BN + Linear [+ L2]."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
d_descriptor: int,
|
||||
return_normalized: bool = False,
|
||||
use_text_film: bool = False,
|
||||
text_film_dim: int = 1024,
|
||||
text_film_hidden: int = 256,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.return_normalized = return_normalized
|
||||
self.use_text_film = use_text_film
|
||||
if use_text_film:
|
||||
self.text_film = TextFiLM(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
|
||||
self.ggem = GGeM(in_channels)
|
||||
self.bn = nn.BatchNorm1d(in_channels, affine=False)
|
||||
self.proj = nn.Linear(in_channels, d_descriptor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
text_emb: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
if self.use_text_film:
|
||||
x = self.text_film(x, text_emb)
|
||||
g = self.ggem(x)
|
||||
g = self.bn(g)
|
||||
g = self.proj(g)
|
||||
if self.return_normalized:
|
||||
g = F.normalize(g, p=2, dim=-1)
|
||||
return g
|
||||
|
||||
|
||||
class UAVHeadV1(nn.Module):
|
||||
"""UAV head: AltitudeFiLM [+ TextFiLM] + GGeM + BN + Linear [+ L2]."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
d_descriptor: int,
|
||||
use_film: bool = True,
|
||||
altitude_norm: float = 500.0,
|
||||
return_normalized: bool = False,
|
||||
use_text_film: bool = False,
|
||||
text_film_dim: int = 1024,
|
||||
text_film_hidden: int = 256,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.return_normalized = return_normalized
|
||||
self.use_film = use_film
|
||||
self.use_text_film = use_text_film
|
||||
if use_film:
|
||||
self.film = AltitudeFiLM(in_channels, altitude_norm=altitude_norm)
|
||||
if use_text_film:
|
||||
self.text_film = TextFiLM(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
|
||||
self.ggem = GGeM(in_channels)
|
||||
self.bn = nn.BatchNorm1d(in_channels, affine=False)
|
||||
self.proj = nn.Linear(in_channels, d_descriptor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
altitude: Optional[torch.Tensor] = None,
|
||||
text_emb: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
if self.use_film:
|
||||
x = self.film(x, altitude)
|
||||
if self.use_text_film:
|
||||
x = self.text_film(x, text_emb)
|
||||
g = self.ggem(x)
|
||||
g = self.bn(g)
|
||||
g = self.proj(g)
|
||||
if self.return_normalized:
|
||||
g = F.normalize(g, p=2, dim=-1)
|
||||
return g
|
||||
93
src/models/sofia_v1/model.py
Normal file
93
src/models/sofia_v1/model.py
Normal file
@@ -0,0 +1,93 @@
|
||||
"""SOFIA v1 model: StripNet+DCNv4 backbone + asymmetric Sat/UAV heads.
|
||||
|
||||
Architecture:
|
||||
img [B,3,256,256] --> StripNetDCN (4 stages) --> [B, C_4, 8, 8]
|
||||
|
|
||||
┌─────────────────────┴────────────────┐
|
||||
▼ ▼
|
||||
SatHeadV1 UAVHeadV1
|
||||
[TextFiLM] AltitudeFiLM
|
||||
GGeM [TextFiLM]
|
||||
BN + Linear GGeM
|
||||
[L2] BN + Linear
|
||||
[L2]
|
||||
|
||||
Heads return un-normalized D-dim descriptors (for downstream gated text fusion).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .config import SOFIAv1Config
|
||||
from .heads import SatHeadV1, UAVHeadV1
|
||||
from .stripnet_model_dcn import VARIANT_MAP, StripNetDCN
|
||||
|
||||
|
||||
class SOFIAv1(nn.Module):
|
||||
"""SOFIA v1: shared StripNet+DCNv4 backbone with asymmetric heads."""
|
||||
|
||||
def __init__(self, cfg: SOFIAv1Config) -> None:
|
||||
super().__init__()
|
||||
if cfg.variant not in VARIANT_MAP:
|
||||
raise ValueError(f"Unknown variant {cfg.variant!r}")
|
||||
self.cfg = cfg
|
||||
|
||||
# Single shared backbone for sat + uav (saves params; v1 stays lightweight).
|
||||
self.backbone: StripNetDCN = VARIANT_MAP[cfg.variant](dcn_variant=cfg.dcn_variant)
|
||||
last_channels = self.backbone.embed_dims[-1]
|
||||
self.feature_channels = last_channels
|
||||
|
||||
self.sat_head = SatHeadV1(
|
||||
in_channels=last_channels,
|
||||
d_descriptor=cfg.d_descriptor,
|
||||
return_normalized=cfg.return_normalized,
|
||||
use_text_film=cfg.use_text_film_sat,
|
||||
text_film_dim=cfg.text_film_dim,
|
||||
text_film_hidden=cfg.text_film_hidden,
|
||||
)
|
||||
self.uav_head = UAVHeadV1(
|
||||
in_channels=last_channels,
|
||||
d_descriptor=cfg.d_descriptor,
|
||||
use_film=cfg.use_film_altitude,
|
||||
altitude_norm=cfg.altitude_norm,
|
||||
return_normalized=cfg.return_normalized,
|
||||
use_text_film=cfg.use_text_film_uav,
|
||||
text_film_dim=cfg.text_film_dim,
|
||||
text_film_hidden=cfg.text_film_hidden,
|
||||
)
|
||||
|
||||
def _extract_last(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Run backbone, return only the deepest feature map [B, C, 8, 8]."""
|
||||
feats: List[torch.Tensor] = self.backbone(x)
|
||||
return feats[-1]
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sat: Optional[torch.Tensor] = None,
|
||||
uav: Optional[torch.Tensor] = None,
|
||||
altitude: Optional[torch.Tensor] = None,
|
||||
text_emb_sat: Optional[torch.Tensor] = None,
|
||||
text_emb_uav: Optional[torch.Tensor] = None,
|
||||
return_features: bool = False,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
result: Dict[str, torch.Tensor] = {}
|
||||
|
||||
if sat is not None:
|
||||
f_sat = self._extract_last(sat)
|
||||
g_sat = self.sat_head(f_sat, text_emb=text_emb_sat)
|
||||
result["g_sat"] = g_sat
|
||||
if return_features:
|
||||
result["features_sat"] = f_sat
|
||||
|
||||
if uav is not None:
|
||||
f_uav = self._extract_last(uav)
|
||||
g_uav = self.uav_head(f_uav, altitude=altitude, text_emb=text_emb_uav)
|
||||
result["g_uav"] = g_uav
|
||||
if return_features:
|
||||
result["features_uav"] = f_uav
|
||||
|
||||
return result
|
||||
57
src/models/sofia_v1/stripnet_blocks_dcn.py
Normal file
57
src/models/sofia_v1/stripnet_blocks_dcn.py
Normal file
@@ -0,0 +1,57 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchvision.ops import DeformConv2d
|
||||
|
||||
class DCNBlock(nn.Module):
|
||||
"""
|
||||
StripNet-style block but uses deformable conv instead of rigid convs.
|
||||
"""
|
||||
def __init__(self, in_ch, out_ch, hidden_ratio=0.25, modulation=True):
|
||||
super().__init__()
|
||||
hidden_ch = max(1, int(out_ch * hidden_ratio))
|
||||
|
||||
# 1x1 reduce
|
||||
self.reduce = nn.Sequential(
|
||||
nn.Conv2d(in_ch, hidden_ch, kernel_size=1, bias=False),
|
||||
nn.BatchNorm2d(hidden_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
# Offset conv (predicts offsets and optional mask)
|
||||
offset_channels = 2 * 3 * 3 if not modulation else 3 * 3 * 3
|
||||
self.offset_conv = nn.Conv2d(hidden_ch, offset_channels,
|
||||
kernel_size=3, padding=1)
|
||||
|
||||
# Deformable conv
|
||||
self.deform = DeformConv2d(hidden_ch, hidden_ch,
|
||||
kernel_size=3, padding=1, bias=False)
|
||||
self.bn = nn.BatchNorm2d(hidden_ch)
|
||||
self.act = nn.ReLU(inplace=True)
|
||||
|
||||
# 1x1 expand
|
||||
self.expand = nn.Sequential(
|
||||
nn.Conv2d(hidden_ch, out_ch, kernel_size=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
)
|
||||
|
||||
self.residual = (in_ch == out_ch)
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
x = self.reduce(x)
|
||||
offset = self.offset_conv(x)
|
||||
|
||||
if offset.shape[1] == 18: # DCNv1
|
||||
x = self.deform(x, offset)
|
||||
else: # DCNv2: last 9 channels are mask
|
||||
o, mask = offset.split([18, 9], dim=1)
|
||||
mask = mask.sigmoid()
|
||||
x = self.deform(x, o, mask)
|
||||
|
||||
x = self.act(self.bn(x))
|
||||
x = self.expand(x)
|
||||
|
||||
if self.residual:
|
||||
x = x + identity
|
||||
return x
|
||||
125
src/models/sofia_v1/stripnet_blocks_dcn_new.py
Normal file
125
src/models/sofia_v1/stripnet_blocks_dcn_new.py
Normal file
@@ -0,0 +1,125 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from DCNv4 import DCNv4 # ← your installed OpenGVLab version
|
||||
import math
|
||||
|
||||
class DCNBlockV4(nn.Module):
|
||||
"""
|
||||
StripNet-style block that uses OpenGVLab DCNv4 instead of torchvision DCNv2.
|
||||
"""
|
||||
|
||||
def __init__(self, in_ch, out_ch, hidden_ratio=0.25,
|
||||
kernel_size=3, stride=1, dilation=1, group=4,
|
||||
offset_scale=1.0, use_bias=False):
|
||||
super().__init__()
|
||||
assert kernel_size in (3, 5, 7)
|
||||
pad = (kernel_size // 2) * dilation
|
||||
|
||||
# Hidden channels — must satisfy (hidden_ch // group) % 16 == 0
|
||||
hidden_ch = max(16, int(out_ch * hidden_ratio))
|
||||
hidden_ch = math.ceil(hidden_ch / 16) * 16
|
||||
# increase until kernel constraint satisfied
|
||||
while (hidden_ch // group) % 16 != 0:
|
||||
hidden_ch += 16
|
||||
|
||||
#print(f"[DCNv4] adjusted hidden_ch={hidden_ch}, group={group}")
|
||||
|
||||
|
||||
self.reduce = nn.Sequential(
|
||||
nn.Conv2d(in_ch, hidden_ch, 1, bias=False),
|
||||
self.make_gn(hidden_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
# DCNv4 core
|
||||
self.dcn = DCNv4(
|
||||
channels=hidden_ch,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
pad=pad,
|
||||
dilation=dilation,
|
||||
group=group,
|
||||
offset_scale=offset_scale,
|
||||
dw_kernel_size=None,
|
||||
center_feature_scale=False,
|
||||
remove_center=False,
|
||||
output_bias=True,
|
||||
without_pointwise=False,
|
||||
)
|
||||
|
||||
|
||||
self.expand = nn.Sequential(
|
||||
nn.Conv2d(hidden_ch, out_ch, 1, bias=False),
|
||||
self.make_gn(out_ch),
|
||||
|
||||
)
|
||||
|
||||
self.residual = (in_ch == out_ch and stride == 1)
|
||||
self.stride = stride
|
||||
|
||||
def make_gn(self, num_channels):
|
||||
num_groups = max(1, num_channels // 16)
|
||||
return nn.GroupNorm(num_groups, num_channels)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Input: [B, C, H, W]
|
||||
Output: [B, C, H', W']
|
||||
"""
|
||||
identity = x
|
||||
|
||||
x = self.reduce(x)
|
||||
B, C, H, W = x.shape
|
||||
|
||||
# Flatten for DCNv4
|
||||
x_seq = x.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
|
||||
|
||||
# --- clamp activations to avoid huge values ---
|
||||
x_seq = torch.clamp(x_seq, -50.0, 50.0)
|
||||
|
||||
x_seq = self.dcn(x_seq)
|
||||
x_seq = torch.nan_to_num(x_seq, nan=0.0, posinf=1e4, neginf=-1e4)
|
||||
|
||||
# --- back to (B, C, H, W) ---
|
||||
x = x_seq.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
# --- expand + residual ---
|
||||
x = self.expand(x)
|
||||
if self.residual:
|
||||
x = x + identity
|
||||
return x
|
||||
|
||||
|
||||
def test_dcnblock_v4():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Using device: {device}")
|
||||
|
||||
block = DCNBlockV4(
|
||||
in_ch=128, # must match input
|
||||
out_ch=128,
|
||||
hidden_ratio=0.25,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
group=5,
|
||||
offset_scale=1.0,
|
||||
).to(device)
|
||||
|
||||
# ✅ internal feature map (not RGB)
|
||||
x = torch.randn(2, 128, 128, 128, device=device, requires_grad=True)
|
||||
|
||||
with torch.cuda.amp.autocast_mode.autocast(enabled=torch.cuda.is_available()):
|
||||
y = block(x)
|
||||
|
||||
print(f"Input shape : {x.shape}")
|
||||
print(f"Output shape : {y.shape}")
|
||||
assert y.shape == x.shape
|
||||
|
||||
loss = y.mean()
|
||||
loss.backward()
|
||||
print("✅ DCNBlockV4 test passed.\n")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_dcnblock_v4()
|
||||
117
src/models/sofia_v1/stripnet_model_dcn.py
Normal file
117
src/models/sofia_v1/stripnet_model_dcn.py
Normal file
@@ -0,0 +1,117 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from src.models.sofia_v1.stripnet_blocks_dcn import DCNBlock
|
||||
from src.models.sofia_v1.stripnet_blocks_dcn_new import DCNBlockV4
|
||||
from src.models.stripnet.model import OverlapPatchEmbed
|
||||
|
||||
|
||||
class StripNetDCN(nn.Module):
|
||||
"""4-stage hierarchical CNN backbone: OverlapPatchEmbed + DCN blocks per stage.
|
||||
|
||||
Output: list of [B, C_i, H_i, W_i] features at each stage. For 256x256 input
|
||||
with default downsampling (4, 2, 2, 2), final stage is 8x8.
|
||||
|
||||
DCN variant:
|
||||
- "v2" (default): torchvision `DeformConv2d` — stable, no memory leaks.
|
||||
- "v4": OpenGVLab DCNv4 — faster on CUDA but has a known C++ extension
|
||||
memory leak (~9 MB per forward call) that causes OOM in long training
|
||||
runs. Only use if you have a patched DCNv4 build.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_chans: int = 3,
|
||||
embed_dims: list[int] = [64, 128, 256, 512],
|
||||
depths: list[int] = [3, 4, 6, 3],
|
||||
dcn_variant: str = "v2",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if dcn_variant not in ("v2", "v4"):
|
||||
raise ValueError(f"dcn_variant must be 'v2' or 'v4', got {dcn_variant!r}")
|
||||
self.num_stages = len(embed_dims)
|
||||
self.embed_dims = embed_dims
|
||||
self.dcn_variant = dcn_variant
|
||||
|
||||
for i in range(self.num_stages):
|
||||
patch_embed = OverlapPatchEmbed(
|
||||
patch_size=7 if i == 0 else 3,
|
||||
stride=4 if i == 0 else 2,
|
||||
in_chans=in_chans if i == 0 else self.embed_dims[i - 1],
|
||||
embed_dim=self.embed_dims[i],
|
||||
)
|
||||
|
||||
block_cls = DCNBlockV4 if dcn_variant == "v4" else DCNBlock
|
||||
block = nn.ModuleList([
|
||||
block_cls(
|
||||
in_ch=self.embed_dims[i],
|
||||
out_ch=self.embed_dims[i],
|
||||
) for _ in range(depths[i])
|
||||
])
|
||||
|
||||
setattr(self, f"patch_embed{i + 1}", patch_embed)
|
||||
setattr(self, f"block{i + 1}", block)
|
||||
|
||||
def forward_features(self, x: torch.Tensor) -> list[torch.Tensor]:
|
||||
outs = []
|
||||
for i in range(self.num_stages):
|
||||
patch_embed = getattr(self, f"patch_embed{i + 1}")
|
||||
block = getattr(self, f"block{i + 1}")
|
||||
x, H, W = patch_embed(x)
|
||||
for blk in block:
|
||||
x = blk(x)
|
||||
outs.append(x)
|
||||
return outs
|
||||
|
||||
def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
|
||||
return self.forward_features(x)
|
||||
|
||||
def get_stripnet_dcn_small(dcn_variant: str = "v2") -> StripNetDCN:
|
||||
return StripNetDCN(
|
||||
in_chans=3,
|
||||
embed_dims=[64, 128, 320, 512],
|
||||
depths=[2, 2, 4, 2],
|
||||
dcn_variant=dcn_variant,
|
||||
)
|
||||
|
||||
|
||||
def get_stripnet_dcn_small_v2(dcn_variant: str = "v2") -> StripNetDCN:
|
||||
return StripNetDCN(
|
||||
in_chans=3,
|
||||
embed_dims=[64, 128, 256, 384],
|
||||
depths=[2, 2, 4, 2],
|
||||
dcn_variant=dcn_variant,
|
||||
)
|
||||
|
||||
|
||||
def get_stripnet_dcn_tiny(dcn_variant: str = "v2") -> StripNetDCN:
|
||||
return StripNetDCN(
|
||||
in_chans=3,
|
||||
embed_dims=[32, 64, 128, 256],
|
||||
depths=[3, 3, 5, 2],
|
||||
dcn_variant=dcn_variant,
|
||||
)
|
||||
|
||||
|
||||
def get_stripnet_dcn_tiny_tiny(dcn_variant: str = "v2") -> StripNetDCN:
|
||||
return StripNetDCN(
|
||||
in_chans=3,
|
||||
embed_dims=[16, 32, 80, 128],
|
||||
depths=[3, 3, 5, 2],
|
||||
dcn_variant=dcn_variant,
|
||||
)
|
||||
|
||||
|
||||
VARIANT_MAP = {
|
||||
"small": get_stripnet_dcn_small,
|
||||
"small_v2": get_stripnet_dcn_small_v2,
|
||||
"tiny": get_stripnet_dcn_tiny,
|
||||
"tiny_tiny": get_stripnet_dcn_tiny_tiny,
|
||||
}
|
||||
|
||||
|
||||
def build_stripnet_dcn(variant: str = "small", dcn_variant: str = "v2") -> StripNetDCN:
|
||||
"""Factory: variant in {tiny_tiny, tiny, small, small_v2}, dcn_variant in {v2, v4}."""
|
||||
if variant not in VARIANT_MAP:
|
||||
raise ValueError(f"Unknown variant {variant!r}. Available: {list(VARIANT_MAP)}")
|
||||
return VARIANT_MAP[variant](dcn_variant=dcn_variant)
|
||||
331
src/models/sofia_v71/README.md
Normal file
331
src/models/sofia_v71/README.md
Normal file
@@ -0,0 +1,331 @@
|
||||
# SOFIA v7.1 — PyTorch implementation
|
||||
|
||||
Reference implementation of **SOFIA v7.1** CVGL student model targeting Jetson Orin NX with 500 MB – 1 GB VRAM after INT8 quantization.
|
||||
|
||||
Full design rationale: see `2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Input (256×256×3, sat or UAV)
|
||||
↓
|
||||
Stem: dual-conv (3→C_mid→C_stem_out), ×2 down
|
||||
↓
|
||||
Stage 1 (shared sat/UAV): StripDCN-lite × d1, ×2 down
|
||||
↓
|
||||
Stage 2 (shared): StripMixConv × d2, ×2 down
|
||||
↓
|
||||
Stage 3 (separate): MambaVision MV5 × d3, ×2 down
|
||||
↓
|
||||
Stage 4 (separate): MambaVision MV1 × d4, ×2 down (→ 8×8)
|
||||
↓
|
||||
Ultra-lite 1×1 Neck → F̃ [B, C_n, 8, 8]
|
||||
↓
|
||||
┌─ Sat-Head: GGeM → BN → Linear → L2 → g_sat
|
||||
└─ UAV-Head: FiLM(altitude) → CHP (polar+FFT+mag) → BN → Linear → L2 → g_uav
|
||||
```
|
||||
|
||||
## Presets
|
||||
|
||||
| Preset | Params | FLOPs | INT8 weights | FP16 weights | Target latency |
|
||||
|--------|-------:|------:|-------------:|-------------:|---------------:|
|
||||
| **M** (default) | ~500 M | ~132 G | ~500 MB | ~1 GB | ~18 ms |
|
||||
| **L** | ~1 B | ~283 G | ~1 GB | ~2 GB | ~20 ms |
|
||||
| Tiny | ~5 M | ~1.4 G | ~5 MB | ~10 MB | ~4 ms |
|
||||
|
||||
*Latency estimates на Jetson Orin NX 8GB (INT8 TRT mixed-precision для DCN/Mamba в FP16).*
|
||||
|
||||
## Dependencies
|
||||
|
||||
Required:
|
||||
- Python ≥ 3.9
|
||||
- PyTorch ≥ 2.0
|
||||
- torchvision ≥ 0.15 (для `deform_conv2d`)
|
||||
|
||||
Optional (для production Mamba speed):
|
||||
- [`mamba_ssm`](https://github.com/state-spaces/mamba) — ускоренные backends:
|
||||
- v1 `selective_scan_fn` для Mamba-1 (~5× vs Python loop)
|
||||
- v2 `Mamba2` модуль для Mamba-2 SSD dual form (~2–8× vs Mamba-1)
|
||||
- [`causal-conv1d`](https://github.com/Dao-AILab/causal-conv1d) — ускоренный 1D conv для Mamba
|
||||
|
||||
## Mamba variants (приоритет по качеству/скорости)
|
||||
|
||||
| Variant | Описание | Speed | Quality | Когда использовать |
|
||||
|---------|----------|:-----:|:-------:|--------------------|
|
||||
| **`mamba2`** (default) | SSD dual form, scalar A per head | **2–8× faster** | best | Main choice если mamba_ssm v2 доступен |
|
||||
| `mamba1` | Original selective scan, diagonal A | ref | +0% vs mamba2 | Legacy / reproducibility / если v2 недоступен |
|
||||
| `efficient_vmamba` | Atrous scan (2 directions, no CUDA kernel) | 2–3× faster | ~−0.3% R@1 | Speed fallback без mamba_ssm |
|
||||
|
||||
### Выбор через config
|
||||
|
||||
```python
|
||||
from code_sofia_v71 import sofia_m_config, SOFIAv71
|
||||
|
||||
cfg = sofia_m_config()
|
||||
|
||||
# Вариант 1: Mamba-2 (default, preferred) — если mamba_ssm v2 доступен
|
||||
cfg.mamba_variant = "mamba2"
|
||||
cfg.mamba_backend = "auto" # falls back to torch if unavailable
|
||||
|
||||
# Вариант 2: EfficientVMamba — speed без зависимости от mamba_ssm
|
||||
cfg.mamba_variant = "efficient_vmamba"
|
||||
|
||||
# Вариант 3: Mamba-1 — legacy / для сравнения
|
||||
cfg.mamba_variant = "mamba1"
|
||||
|
||||
model = SOFIAv71(cfg)
|
||||
```
|
||||
|
||||
### Проверка доступности backends
|
||||
|
||||
```python
|
||||
from code_sofia_v71 import is_mamba_ssm_available, is_mamba2_available
|
||||
|
||||
print(f"Mamba-1 CUDA: {is_mamba_ssm_available()}")
|
||||
print(f"Mamba-2 CUDA: {is_mamba2_available()}")
|
||||
```
|
||||
|
||||
### Несовместимость параметров между variants
|
||||
|
||||
Mamba-1 и Mamba-2 используют **разную параметризацию $A$** (diagonal-per-channel vs scalar-per-head) — state_dict **НЕ совместим**. EfficientVMamba имеет N независимых scanners каждый с собственными параметрами Mamba-1.
|
||||
|
||||
**Следствие:** нельзя train в `mamba1` и load в `mamba2` checkpoint. Выбор variant нужно зафиксировать до training.
|
||||
|
||||
## Quick start
|
||||
|
||||
```python
|
||||
import torch
|
||||
from code_sofia_v71 import build_sofia
|
||||
|
||||
# Default = SOFIA-M (500 MB INT8 target)
|
||||
model = build_sofia("M")
|
||||
model.eval()
|
||||
|
||||
sat = torch.randn(2, 3, 256, 256) # satellite image
|
||||
uav = torch.randn(2, 3, 256, 256) # UAV image
|
||||
altitude = torch.tensor([120.0, 450.0]) # meters
|
||||
|
||||
out = model(sat=sat, uav=uav, altitude=altitude)
|
||||
g_sat = out["g_sat"] # [2, 512] L2-normalized
|
||||
g_uav = out["g_uav"] # [2, 512] L2-normalized
|
||||
|
||||
# Retrieval similarity
|
||||
similarity = (g_sat * g_uav).sum(dim=-1) # cosine
|
||||
```
|
||||
|
||||
## Verification
|
||||
|
||||
```bash
|
||||
# Smoke test (Tiny preset, CPU, fast)
|
||||
python -m sofia_v71.verify
|
||||
|
||||
# Full SOFIA-M check with latency benchmark (GPU)
|
||||
python -m sofia_v71.verify --preset M --device cuda --benchmark
|
||||
|
||||
# Maximum scale
|
||||
python -m sofia_v71.verify --preset L --device cuda
|
||||
```
|
||||
|
||||
Expected output for SOFIA-M:
|
||||
```
|
||||
Total params: ~500 M
|
||||
...
|
||||
FP16 weights: ~1000 MB (~0.98 GB)
|
||||
INT8 weights: ~500 MB (~0.49 GB)
|
||||
...
|
||||
out[g_sat]: (1, 512)
|
||||
out[g_uav]: (1, 512)
|
||||
```
|
||||
|
||||
## Key novel components
|
||||
|
||||
### `layers.py`
|
||||
- **`GGeM`** — per-channel learnable exponent Generalized Mean pooling (F11)
|
||||
- **`CircularHarmonicPool`** — formally SO(2)-invariant UAV pool via polar → 1D FFT → magnitude (NOVEL NH2)
|
||||
- **`AltitudeFiLM`** — telemetry-conditioned modulation (NOVEL NH4)
|
||||
- **`RoPE2D`** — 2D rotary positional embedding
|
||||
|
||||
### `blocks.py`
|
||||
- **`StripDCNLiteBlock`** — Strip DW MBConv с DCN offset на одной оси (NOVEL)
|
||||
- **`StripMixConvBlock`** — Strip + MixConv (3/5/7 kernels) для multi-scale
|
||||
- **`MambaVisionBlock`** — Mamba ∥ MHSA [∥ Strip] + FFN (MV5 variant — NOVEL 3-way)
|
||||
- **`SimpleMambaBlock`** — reference pure-PyTorch Mamba-1 (replace with `mamba_ssm` в production)
|
||||
|
||||
### `model.py`
|
||||
- **`SOFIAv71`** — full model with optional weight-sharing
|
||||
- **`SatHead` vs `UAVHead`** — asymmetric physics-motivated design (NOVEL NH1)
|
||||
- **`RingAuxHead`** — LPN Square-Ring training-only aux
|
||||
|
||||
## KD taps
|
||||
|
||||
Backbone forward returns features at stages s0/f1/f2/f3/f4. Использовать для
|
||||
hierarchical knowledge distillation:
|
||||
|
||||
```python
|
||||
out = model(sat=sat, uav=uav, return_features=True)
|
||||
f3_sat = out["features_sat"]["f3"] # для teacher feature alignment
|
||||
```
|
||||
|
||||
## Production notes
|
||||
|
||||
### Mamba backend selection
|
||||
|
||||
Default `mamba_variant="mamba2"` с `backend="auto"`:
|
||||
- Use `Mamba2` from `mamba_ssm.modules.mamba2` если установлено (CUDA, fast)
|
||||
- Иначе fallback на pure-PyTorch simplified SSD scan (slow, работает)
|
||||
|
||||
Для Mamba-1 legacy путь — `mamba_variant="mamba1"`, backend так же резолвится.
|
||||
|
||||
Для EfficientVMamba — pure PyTorch, не нуждается в mamba_ssm.
|
||||
|
||||
## Quantization (`quant.py`)
|
||||
|
||||
Reusable utilities для INT8 PTQ/QAT в SOFIA.
|
||||
|
||||
### `OffsetClampSTE` — DCN-M2
|
||||
|
||||
Hard clamp DCN offsets к `[-k, +k]` со STE backward. Уже подключён в `StripDCNLiteBlock` через флаг `use_offset_clamp_ste=True` (default).
|
||||
|
||||
```python
|
||||
from code_sofia_v71 import OffsetClampSTE, offset_clamp_ste
|
||||
|
||||
# Module form
|
||||
clamp = OffsetClampSTE(kernel_size=7)
|
||||
clamped = clamp(offsets)
|
||||
|
||||
# Functional form
|
||||
clamped = offset_clamp_ste(offsets, kernel_size=7)
|
||||
```
|
||||
|
||||
Backward — identity (gradient проходит насквозь). Можно ставить с epoch 1, не только в QAT.
|
||||
|
||||
### `KScaledFakeQuant` — k-scaled fake quantization
|
||||
|
||||
Multi-bin fake quant для long-tail distributions (Mamba Δ, y).
|
||||
|
||||
```python
|
||||
from code_sofia_v71 import KScaledFakeQuant
|
||||
|
||||
fq = KScaledFakeQuant(num_bins=3)
|
||||
|
||||
# Calibration
|
||||
fq.start_calibration()
|
||||
with torch.no_grad():
|
||||
for batch in cal_loader:
|
||||
_ = model(batch)
|
||||
fq.finalize_calibration(percentile=99.9)
|
||||
# fq теперь в PTQ-режиме
|
||||
```
|
||||
|
||||
### `KScaledMamba2Block` — drop-in для Mamba2Block
|
||||
|
||||
Подкласс `Mamba2Block` с k-scaled fake-quant узлами на критических путях
|
||||
(`x_main`, `delta`, `y`). Внутреннее состояние scan'а `h_t` остаётся в
|
||||
model dtype (R5 reparam principle).
|
||||
|
||||
```python
|
||||
from code_sofia_v71 import KScaledMamba2Block
|
||||
|
||||
mamba = KScaledMamba2Block(
|
||||
channels=192,
|
||||
d_state=64,
|
||||
headdim=64,
|
||||
num_bins=3,
|
||||
targets=("x_main", "delta", "y"), # subset
|
||||
)
|
||||
|
||||
mamba.start_calibration()
|
||||
# ... run calibration data ...
|
||||
mamba.finalize_calibration(percentile=99.9)
|
||||
```
|
||||
|
||||
**Constraint:** только `backend='torch'` поддерживается (mamba_ssm CUDA kernel
|
||||
не имеет hooks для k-scaled fake-quant). Для full INT8 deploy на TRT нужен
|
||||
custom plugin — отдельный deploy concern.
|
||||
|
||||
### Model-wide helpers
|
||||
|
||||
```python
|
||||
from code_sofia_v71 import start_calibration, finalize_calibration, set_quant_enabled
|
||||
|
||||
# Начать калибровку всех KScaledFakeQuant и KScaledMamba2Block в модели
|
||||
start_calibration(model)
|
||||
with torch.no_grad():
|
||||
for batch in cal_loader:
|
||||
_ = model(batch)
|
||||
finalize_calibration(model, percentile=99.9)
|
||||
|
||||
# Toggle on/off для FP-vs-INT8 ablation
|
||||
set_quant_enabled(model, False) # FP forward
|
||||
y_fp = model(x)
|
||||
set_quant_enabled(model, True)
|
||||
y_q = model(x)
|
||||
```
|
||||
|
||||
### Smoke test
|
||||
|
||||
```bash
|
||||
python -m sofia_v71.quant
|
||||
```
|
||||
|
||||
Прогонит unit-test на:
|
||||
- DCN-M2 clamp (gradient pass-through)
|
||||
- KScaledFakeQuant calibration + quant error
|
||||
- KScaledMamba2Block FP-vs-INT8 diff
|
||||
|
||||
### DCN INT8 QAT
|
||||
StripDCN использует `torchvision.ops.deform_conv2d`. Для INT8 deploy на TRT:
|
||||
- Применить QAT с DCN-M1..M4 modifications (см. HYP Phase 7)
|
||||
- Offset predictor: per-channel scale, offset clamping `[-k, +k]`
|
||||
- Mask: FP16 micro-block внутри INT8 graph
|
||||
|
||||
### TensorRT export
|
||||
```python
|
||||
import torch
|
||||
model.eval()
|
||||
model_fuse = model # apply fuse_reparam passes separately
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(sat, uav, altitude),
|
||||
"sofia_v71.onnx",
|
||||
input_names=["sat", "uav", "altitude"],
|
||||
output_names=["g_sat", "g_uav"],
|
||||
opset_version=17,
|
||||
)
|
||||
# Then:
|
||||
# trtexec --onnx=sofia_v71.onnx --int8 --fp16 --saveEngine=sofia.plan
|
||||
```
|
||||
|
||||
## Training
|
||||
|
||||
See `HYP_SOFIA_v7_UltraDeep_дизайн.md` Phase 8 for full training recipe:
|
||||
- Loss: InfoNCE (Sample4Geo mining) + Ring aux + (opt) cross-view consistency
|
||||
- Curriculum: PALW sigmoid warmup, 60 epochs, 4 phases
|
||||
- Optimizer: AdamW, LR 3e-4 cosine, wd 0.05
|
||||
- Temperature: τ 0.1 → 0.01 cosine decay
|
||||
|
||||
## File layout
|
||||
|
||||
```
|
||||
code_sofia_v71/
|
||||
├── __init__.py — public exports
|
||||
├── config.py — SOFIAConfig + presets (M/L/Tiny)
|
||||
├── layers.py — GGeM, CHP, FiLM, RoPE2D, SE, LayerNorm2d
|
||||
├── blocks.py — StripDCN, StripMixConv, Mamba, MambaVision, Downsample
|
||||
├── model.py — Stem, Backbone, Neck, Heads, SOFIAv71
|
||||
├── verify.py — parameter counter + benchmark script
|
||||
└── README.md — this file
|
||||
```
|
||||
|
||||
## References to design doc
|
||||
|
||||
Все architectural decisions обоснованы в
|
||||
`2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`:
|
||||
- Phase 1: requirements R1–R13
|
||||
- Phase 2': DCN / Strip / MambaVision operator catalog
|
||||
- Phase 3': backbone design with 4 candidates (E/F/G)
|
||||
- Phase 4'': CVGL-Aware Head v7.1-α (Asymmetric + CHP + FiLM)
|
||||
- Phase 5': ablation matrix
|
||||
- Phase 6: MambaVision MV5 operational details
|
||||
- Phase 7: StripDCN QAT strategy
|
||||
- Phase 8: training pipeline
|
||||
126
src/models/sofia_v71/__init__.py
Normal file
126
src/models/sofia_v71/__init__.py
Normal file
@@ -0,0 +1,126 @@
|
||||
"""SOFIA v7.1 — student model for cross-view geo-localization (CVGL).
|
||||
|
||||
Architecture:
|
||||
stem → stage1 (StripDCN-lite) → stage2 (StripMixConv)
|
||||
→ stage3 (MambaVision MV5) → stage4 (MambaVision MV1)
|
||||
→ 1×1 neck → asymmetric Sat/UAV heads
|
||||
|
||||
Key novel components:
|
||||
- StripDCN-lite: Strip DW with adaptive offset on one axis
|
||||
- MambaVision MV5: Mamba ∥ MHSA ∥ Strip 3-way parallel
|
||||
- CircularHarmonicPool: formally SO(2)-invariant UAV head
|
||||
- AltitudeFiLM: telemetry-aware conditioning
|
||||
|
||||
Presets (see config.py):
|
||||
- SOFIA-M (~500 M params, 500 MB INT8) — default
|
||||
- SOFIA-L (~1 B params, 1 GB INT8) — max scale
|
||||
- SOFIA-Tiny (~5 M) — reference from original v7.1 spec
|
||||
|
||||
Quick start:
|
||||
from sofia_v71 import build_sofia
|
||||
model = build_sofia("M")
|
||||
out = model(sat=sat_tensor, uav=uav_tensor, altitude=alt_tensor)
|
||||
g_sat, g_uav = out["g_sat"], out["g_uav"]
|
||||
|
||||
See README.md and HYP_SOFIA_v7_UltraDeep_дизайн.md for full design rationale.
|
||||
"""
|
||||
|
||||
from .config import (
|
||||
SOFIAConfig,
|
||||
sofia_l_config,
|
||||
sofia_m_config,
|
||||
sofia_tiny_config,
|
||||
DEFAULT_CONFIG,
|
||||
)
|
||||
from .model import (
|
||||
SOFIAv71,
|
||||
build_sofia,
|
||||
Backbone,
|
||||
Stem,
|
||||
UltraLiteNeck,
|
||||
SatHead,
|
||||
UAVHead,
|
||||
RingAuxHead,
|
||||
)
|
||||
from .layers import (
|
||||
GGeM,
|
||||
CircularHarmonicPool,
|
||||
AltitudeFiLM,
|
||||
TextFiLM,
|
||||
RoPE2D,
|
||||
SqueezeExcite,
|
||||
LayerNorm2d,
|
||||
)
|
||||
from .blocks import (
|
||||
StripDCNLiteBlock,
|
||||
StripMixConvBlock,
|
||||
SimpleMambaBlock,
|
||||
Mamba2Block,
|
||||
EfficientVMambaBlock,
|
||||
MambaVisionBlock,
|
||||
EVSSBridge,
|
||||
Downsample,
|
||||
build_mamba_block,
|
||||
is_mamba_ssm_available,
|
||||
is_mamba2_available,
|
||||
)
|
||||
from .quant import (
|
||||
# DCN-M2
|
||||
OffsetClampSTE,
|
||||
offset_clamp_ste,
|
||||
# k-scaled
|
||||
KScaledFakeQuant,
|
||||
KScaledMamba2Block,
|
||||
# model-wide helpers
|
||||
start_calibration,
|
||||
finalize_calibration,
|
||||
set_quant_enabled,
|
||||
)
|
||||
|
||||
__version__ = "0.1.0"
|
||||
__all__ = [
|
||||
# Config
|
||||
"SOFIAConfig",
|
||||
"sofia_m_config",
|
||||
"sofia_l_config",
|
||||
"sofia_tiny_config",
|
||||
"DEFAULT_CONFIG",
|
||||
# Top-level
|
||||
"SOFIAv71",
|
||||
"build_sofia",
|
||||
# Model parts
|
||||
"Backbone",
|
||||
"Stem",
|
||||
"UltraLiteNeck",
|
||||
"SatHead",
|
||||
"UAVHead",
|
||||
"RingAuxHead",
|
||||
# Layers
|
||||
"GGeM",
|
||||
"CircularHarmonicPool",
|
||||
"AltitudeFiLM",
|
||||
"TextFiLM",
|
||||
"RoPE2D",
|
||||
"SqueezeExcite",
|
||||
"LayerNorm2d",
|
||||
# Blocks
|
||||
"StripDCNLiteBlock",
|
||||
"StripMixConvBlock",
|
||||
"SimpleMambaBlock",
|
||||
"Mamba2Block",
|
||||
"EfficientVMambaBlock",
|
||||
"MambaVisionBlock",
|
||||
"EVSSBridge",
|
||||
"Downsample",
|
||||
"build_mamba_block",
|
||||
"is_mamba_ssm_available",
|
||||
"is_mamba2_available",
|
||||
# Quantization
|
||||
"OffsetClampSTE",
|
||||
"offset_clamp_ste",
|
||||
"KScaledFakeQuant",
|
||||
"KScaledMamba2Block",
|
||||
"start_calibration",
|
||||
"finalize_calibration",
|
||||
"set_quant_enabled",
|
||||
]
|
||||
1186
src/models/sofia_v71/blocks.py
Normal file
1186
src/models/sofia_v71/blocks.py
Normal file
File diff suppressed because it is too large
Load Diff
210
src/models/sofia_v71/config.py
Normal file
210
src/models/sofia_v71/config.py
Normal file
@@ -0,0 +1,210 @@
|
||||
"""SOFIA v7.1 configuration system.
|
||||
|
||||
Two scale presets targeting Jetson Orin NX INT8 deployment:
|
||||
|
||||
| Preset | Params | INT8 size | FP16 size | Target latency |
|
||||
|----------|---------:|----------:|----------:|---------------:|
|
||||
| SOFIA-M | ~500 M | ~500 MB | ~1 GB | ~18 ms |
|
||||
| SOFIA-L | ~1 B | ~1 GB | ~2 GB | ~20 ms |
|
||||
|
||||
Based on HYP_SOFIA_v7_UltraDeep_дизайн.md (Phases 2'–5' + Phase 4''
|
||||
CVGL-Aware Head + revision to ultra-lite 1x1 neck).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Literal, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class SOFIAConfig:
|
||||
"""Configuration for SOFIA v7.1 student architecture.
|
||||
|
||||
Controls backbone width/depth, neck, and CVGL-Aware Head.
|
||||
"""
|
||||
|
||||
# -------- Input --------
|
||||
input_size: int = 256
|
||||
in_channels: int = 3
|
||||
|
||||
# -------- Stem (dual-conv FastViT-style) --------
|
||||
stem_mid: int = 64 # intermediate channels
|
||||
stem_out: int = 128 # output channels into stage 1
|
||||
|
||||
# -------- Backbone dimensions per stage s1..s4 --------
|
||||
embed_dims: List[int] = field(
|
||||
default_factory=lambda: [256, 512, 1280, 1536]
|
||||
)
|
||||
depths: List[int] = field(default_factory=lambda: [3, 4, 15, 3])
|
||||
|
||||
# -------- Stage 1–2 block parameters --------
|
||||
mbconv_expand: int = 4
|
||||
se_ratio: int = 16
|
||||
strip_kernel_s1: int = 7 # strip DW kernel size stage 1
|
||||
strip_kernel_s2: int = 5 # strip DW kernel size stage 2
|
||||
mix_kernels: List[int] = field(
|
||||
default_factory=lambda: [3, 5, 7]
|
||||
) # MixConv DW kernel sizes stage 2
|
||||
use_dcn_strip: bool = True # adaptive offset on horizontal strip
|
||||
|
||||
# -------- Stage 3–4 (MambaVision) --------
|
||||
mamba_d_state: int = 16
|
||||
mamba_dt_rank: Optional[int] = None # auto = max(1, C // 16)
|
||||
mamba_backend: Literal["auto", "torch", "mamba_ssm"] = "auto"
|
||||
# "auto" uses mamba_ssm if importable, else torch fallback
|
||||
|
||||
# Mamba variant — one of:
|
||||
# "mamba2" (preferred, SSD dual form, 2-8x faster)
|
||||
# "mamba1" (original selective scan, mature)
|
||||
# "efficient_vmamba" (speed fallback, atrous scan ~2-3x speedup)
|
||||
mamba_variant: Literal["mamba1", "mamba2", "efficient_vmamba"] = "mamba2"
|
||||
# Per-variant tunables passed through to factory
|
||||
mamba_extra_kwargs: dict = field(default_factory=lambda: {
|
||||
"d_state_mamba2": 64, # Mamba-2 typically uses N=64 (not 16)
|
||||
"headdim": 64, # Mamba-2 head dim
|
||||
"expand": 2, # Mamba-2 inner expansion factor
|
||||
"d_conv": 4, # Mamba-2 local conv kernel
|
||||
"n_directions": 2, # EfficientVMamba: 2 or 4 atrous directions
|
||||
})
|
||||
|
||||
num_heads_s3: int = 8
|
||||
num_heads_s4: int = 8
|
||||
use_strip_branch_s3: bool = True # MV5 with Strip branch
|
||||
use_strip_branch_s4: bool = False # MV1 without Strip branch
|
||||
ffn_expand: int = 4
|
||||
|
||||
# -------- EVSS-style bridge (B6-inspired, opt-in) --------
|
||||
# When True, inserts a within-resolution dual-path refinement block
|
||||
# right after each downsample to stage 3 and stage 4. Smooths semantic
|
||||
# gap between heterogeneous stages (DCN → MambaVision). Adds ~165K params
|
||||
# and ~30 MMAC per insertion at C=192, HW=256.
|
||||
use_evss_bridge: bool = False
|
||||
evss_bridge_locations: List[str] = field(
|
||||
default_factory=lambda: ["pre_stage3"] # subset of {pre_stage3, pre_stage4}
|
||||
)
|
||||
|
||||
# -------- Neck (ultra-lite 1x1 projection) --------
|
||||
neck_channels: int = 192 # C_n, output channels from neck
|
||||
|
||||
# -------- CVGL-Aware Head v7.1-α --------
|
||||
d_descriptor: int = 512 # global descriptor dimensionality
|
||||
use_asymmetric_heads: bool = True # Sat vs UAV different heads
|
||||
chp_rings: int = 8
|
||||
chp_angles: int = 16
|
||||
chp_harmonics: int = 4
|
||||
use_film_altitude: bool = True
|
||||
altitude_norm: float = 500.0 # divides altitude in meters
|
||||
ring_count: int = 4 # LPN rings auxiliary
|
||||
use_ring_aux: bool = True # training-only ring aux branch
|
||||
|
||||
# -------- Text fusion (extension for caption-conditioned heads) --------
|
||||
return_normalized: bool = True # if False, heads return pre-L2 features (for late gated fusion)
|
||||
use_text_film_sat: bool = False # text-FiLM modulation in SatHead before GGeM
|
||||
use_text_film_uav: bool = False # text-FiLM modulation in UAVHead alongside altitude FiLM
|
||||
text_film_dim: int = 1024 # text embedding dim feeding FiLM (matches TextFusionMLP out_dim)
|
||||
text_film_hidden: int = 256
|
||||
|
||||
# -------- Weight-sharing --------
|
||||
share_stages_1_2: bool = True # sat ↔ UAV shared weights stages 1-2
|
||||
|
||||
# -------- KD taps (enable for future teacher KD) --------
|
||||
enable_kd_taps: bool = True
|
||||
|
||||
# -------- Deployment hints --------
|
||||
precision: Literal["fp32", "fp16", "int8_mixed"] = "fp16"
|
||||
|
||||
def validate(self) -> None:
|
||||
assert len(self.embed_dims) == 4, (
|
||||
f"embed_dims must have 4 entries, got {len(self.embed_dims)}"
|
||||
)
|
||||
assert len(self.depths) == 4, (
|
||||
f"depths must have 4 entries, got {len(self.depths)}"
|
||||
)
|
||||
assert self.input_size % 32 == 0, (
|
||||
f"input_size must be divisible by 32 (4 downsamples × stem), "
|
||||
f"got {self.input_size}"
|
||||
)
|
||||
|
||||
def summary(self) -> str:
|
||||
return (
|
||||
f"SOFIAConfig(stem={self.stem_mid}/{self.stem_out}, "
|
||||
f"dims={self.embed_dims}, depths={self.depths}, "
|
||||
f"neck={self.neck_channels}, d={self.d_descriptor}, "
|
||||
f"precision={self.precision})"
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Scale Presets
|
||||
# ============================================================
|
||||
|
||||
def sofia_m_config() -> SOFIAConfig:
|
||||
"""SOFIA-M: ~500 M params target (~500 MB INT8, ~1 GB FP16).
|
||||
|
||||
Fits in 500 MB VRAM after INT8 quantization on Jetson Orin NX.
|
||||
Expected latency: ~18 ms.
|
||||
"""
|
||||
return SOFIAConfig(
|
||||
stem_mid=64,
|
||||
stem_out=128,
|
||||
embed_dims=[256, 512, 1280, 1536],
|
||||
depths=[3, 4, 15, 3],
|
||||
neck_channels=192,
|
||||
d_descriptor=512,
|
||||
)
|
||||
|
||||
|
||||
def sofia_l_config() -> SOFIAConfig:
|
||||
"""SOFIA-L: ~1 B params target (~1 GB INT8, ~2 GB FP16).
|
||||
|
||||
Fits in 1 GB VRAM after INT8 quantization on Jetson Orin NX.
|
||||
Expected latency: ~20 ms.
|
||||
"""
|
||||
return SOFIAConfig(
|
||||
stem_mid=64,
|
||||
stem_out=128,
|
||||
embed_dims=[256, 512, 1536, 2048],
|
||||
depths=[3, 4, 20, 3],
|
||||
neck_channels=256,
|
||||
d_descriptor=1024,
|
||||
)
|
||||
|
||||
|
||||
def sofia_tiny_config() -> SOFIAConfig:
|
||||
"""SOFIA-Tiny: ~5 M params (matches original v7.1 spec).
|
||||
|
||||
Reference for research comparisons. Not optimized for 500 MB INT8 target.
|
||||
`num_heads_*` is set to 4 so `head_dim` (channels // heads) is divisible
|
||||
by 4 — required by RoPE 2D in `MambaVisionBlock` (s3: 176/4=44, s4: 224/4=56).
|
||||
`mamba_extra_kwargs.headdim=16` because Mamba-2 requires channels % headdim == 0;
|
||||
176 and 224 are not divisible by the default 64.
|
||||
"""
|
||||
return SOFIAConfig(
|
||||
stem_mid=16,
|
||||
stem_out=32,
|
||||
embed_dims=[48, 96, 176, 224],
|
||||
depths=[2, 3, 4, 2],
|
||||
num_heads_s3=4,
|
||||
num_heads_s4=4,
|
||||
neck_channels=128,
|
||||
d_descriptor=512,
|
||||
mamba_extra_kwargs={
|
||||
"d_state_mamba2": 64,
|
||||
"headdim": 16,
|
||||
"expand": 2,
|
||||
"d_conv": 4,
|
||||
"n_directions": 2,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Default preset
|
||||
DEFAULT_CONFIG = sofia_m_config
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
for name, fn in [("M", sofia_m_config), ("L", sofia_l_config), ("Tiny", sofia_tiny_config)]:
|
||||
cfg = fn()
|
||||
cfg.validate()
|
||||
print(f"SOFIA-{name}: {cfg.summary()}")
|
||||
398
src/models/sofia_v71/layers.py
Normal file
398
src/models/sofia_v71/layers.py
Normal file
@@ -0,0 +1,398 @@
|
||||
"""SOFIA v7.1 custom layers.
|
||||
|
||||
Includes:
|
||||
- GGeM: Generalized Mean Pooling with per-channel learnable exponent (F11)
|
||||
- CircularHarmonicPool: Formally SO(2)-invariant pooling via polar + FFT magnitude (NH2 novel)
|
||||
- AltitudeFiLM: FiLM conditioning on UAV altitude (NH4 novel)
|
||||
- RoPE2D: 2D Rotary Position Embedding for attention
|
||||
- SqueezeExcite: standard SE block
|
||||
- LayerNorm2d: channel-last LN wrapper for 2D features
|
||||
|
||||
All rotation-invariance and FiLM modules are NOVEL contributions of SOFIA v7.1-α.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# ============================================================
|
||||
# GGeM: Generalized Mean Pooling (F11)
|
||||
# ============================================================
|
||||
|
||||
class GGeM(nn.Module):
|
||||
"""Per-channel learnable Generalized Mean pooling.
|
||||
|
||||
Formula:
|
||||
GGeM(F)_c = (1/HW · Σ F_{c,h,w}^{p_c})^{1/p_c}
|
||||
p_c = softplus(p_hat_c) ∈ (0, ∞)
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int, init_p: float = 3.0, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
# softplus^{-1}(init_p) = log(exp(init_p) - 1)
|
||||
hat_init = math.log(math.exp(init_p) - 1.0)
|
||||
self.hat_p = nn.Parameter(torch.full((channels,), hat_init))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""x: [B, C, H, W] -> [B, C]"""
|
||||
p = F.softplus(self.hat_p).view(1, -1, 1, 1) # [1, C, 1, 1]
|
||||
x_clamped = x.clamp(min=self.eps)
|
||||
x_pow = x_clamped.pow(p)
|
||||
x_mean = x_pow.mean(dim=(2, 3), keepdim=True)
|
||||
out = x_mean.pow(1.0 / p)
|
||||
return out.flatten(1)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# CHP: Circular Harmonic Pool (NH2 novel — formally SO(2)-invariant)
|
||||
# ============================================================
|
||||
|
||||
class CircularHarmonicPool(nn.Module):
|
||||
"""Formally SO(2) rotation-invariant pooling.
|
||||
|
||||
Algorithm:
|
||||
1. Sample input feature map at polar grid (r, θ) via bilinear grid_sample
|
||||
2. Apply 1D real FFT along θ-axis
|
||||
3. Keep magnitudes of first N harmonics (invariant to shift = rotation)
|
||||
4. GGeM pool over rings r (per-channel-per-harmonic)
|
||||
5. Flatten to descriptor [B, C * N]
|
||||
|
||||
Output is theoretically invariant to input rotation of any angle.
|
||||
|
||||
See HYP Phase 4'' Section 4''.1 NH2 and Section 4''.5 for formal proof.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
rings: int = 8,
|
||||
angles: int = 16,
|
||||
harmonics: int = 4,
|
||||
r_min: float = 0.1,
|
||||
r_max: float = 1.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
assert harmonics <= angles // 2 + 1, (
|
||||
f"harmonics {harmonics} cannot exceed angles//2+1 = {angles // 2 + 1}"
|
||||
)
|
||||
|
||||
self.channels = channels
|
||||
self.rings = rings
|
||||
self.angles = angles
|
||||
self.harmonics = harmonics
|
||||
|
||||
# GGeM over rings (per channel × per harmonic)
|
||||
self.ggem = GGeM(channels * harmonics, init_p=3.0)
|
||||
|
||||
# Precompute polar grid in normalized [-1, 1] coords for grid_sample
|
||||
grid = self._make_polar_grid(rings, angles, r_min, r_max) # [R, T, 2]
|
||||
self.register_buffer("polar_grid", grid, persistent=False)
|
||||
|
||||
@staticmethod
|
||||
def _make_polar_grid(R: int, T: int, r_min: float, r_max: float) -> torch.Tensor:
|
||||
r_values = torch.linspace(r_min, r_max, R) # [R]
|
||||
theta_values = torch.linspace(0.0, 2 * math.pi, T + 1)[:-1] # [T]
|
||||
r_grid, theta_grid = torch.meshgrid(r_values, theta_values, indexing="ij") # [R, T]
|
||||
x = r_grid * torch.cos(theta_grid)
|
||||
y = r_grid * torch.sin(theta_grid)
|
||||
grid = torch.stack([x, y], dim=-1) # [R, T, 2]
|
||||
return grid
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [B, C, H, W] feature map
|
||||
Returns:
|
||||
descriptor: [B, C * harmonics]
|
||||
"""
|
||||
B, C, H, W = x.shape
|
||||
assert C == self.channels, (
|
||||
f"Input channels {C} != expected {self.channels}"
|
||||
)
|
||||
|
||||
# 1. Polar sampling
|
||||
# grid: [B, R, T, 2]
|
||||
grid = self.polar_grid.unsqueeze(0).expand(B, -1, -1, -1)
|
||||
polar = F.grid_sample(
|
||||
x, grid,
|
||||
mode="bilinear",
|
||||
padding_mode="zeros",
|
||||
align_corners=True,
|
||||
) # [B, C, R, T]
|
||||
|
||||
# 2. 1D real FFT along angular axis
|
||||
polar_fft = torch.fft.rfft(polar, dim=-1) # [B, C, R, T//2+1]
|
||||
polar_fft = polar_fft[..., : self.harmonics] # [B, C, R, N]
|
||||
|
||||
# 3. Magnitude (rotation invariant)
|
||||
magnitude = polar_fft.abs() # [B, C, R, N]
|
||||
|
||||
# 4. Reshape for GGeM: treat (C, N) as combined channel dim, rings as spatial
|
||||
# Shape: [B, C*N, R, 1]
|
||||
magnitude_reshaped = (
|
||||
magnitude
|
||||
.permute(0, 1, 3, 2) # [B, C, N, R]
|
||||
.reshape(B, C * self.harmonics, self.rings, 1)
|
||||
)
|
||||
|
||||
# 5. GGeM pool over rings (H=R, W=1)
|
||||
descriptor = self.ggem(magnitude_reshaped) # [B, C*N]
|
||||
|
||||
return descriptor
|
||||
|
||||
|
||||
# ============================================================
|
||||
# FiLM: Altitude-conditioned modulation (NH4 novel)
|
||||
# ============================================================
|
||||
|
||||
class AltitudeFiLM(nn.Module):
|
||||
"""FiLM modulation conditioned on scalar altitude.
|
||||
|
||||
F' = γ(h) · F + β(h) where γ,β ∈ R^C are produced by MLP.
|
||||
|
||||
At altitude=None, produces identity (γ=1, β=0) via zero-init of final layer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
hidden_dim: int = 64,
|
||||
altitude_norm: float = 500.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.altitude_norm = altitude_norm
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(1, hidden_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(hidden_dim, 2 * channels),
|
||||
)
|
||||
# Zero-init last layer → initial γ=0 before residual, β=0
|
||||
nn.init.zeros_(self.mlp[-1].weight)
|
||||
nn.init.zeros_(self.mlp[-1].bias)
|
||||
|
||||
def forward(self, x: torch.Tensor, altitude: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [B, C, H, W]
|
||||
altitude: [B] or [B, 1] scalar altitude in meters (or None for neutral)
|
||||
Returns:
|
||||
[B, C, H, W]
|
||||
"""
|
||||
B = x.shape[0]
|
||||
if altitude is None:
|
||||
altitude = torch.zeros(B, 1, device=x.device, dtype=x.dtype)
|
||||
elif altitude.dim() == 1:
|
||||
altitude = altitude.unsqueeze(-1)
|
||||
|
||||
h_norm = altitude.to(x.dtype) / self.altitude_norm
|
||||
gamma_beta = self.mlp(h_norm) # [B, 2C]
|
||||
gamma, beta = gamma_beta.chunk(2, dim=-1)
|
||||
# Residual form: γ = 1 + delta_γ (starts at identity)
|
||||
gamma = gamma.view(B, self.channels, 1, 1) + 1.0
|
||||
beta = beta.view(B, self.channels, 1, 1)
|
||||
return gamma * x + beta
|
||||
|
||||
|
||||
# ============================================================
|
||||
# TextFiLM: text-conditioned modulation (extension for caption fusion)
|
||||
# ============================================================
|
||||
|
||||
class TextFiLM(nn.Module):
|
||||
"""FiLM modulation conditioned on a text embedding.
|
||||
|
||||
F' = γ(z) · F + β(z) where γ,β ∈ R^C are produced by an MLP
|
||||
from z ∈ R^{D_txt}. Identity at init via zero-init of last layer
|
||||
(so γ=1, β=0 before training shifts the residual).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
text_dim: int = 1024,
|
||||
hidden_dim: int = 256,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.text_dim = text_dim
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(text_dim, hidden_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(hidden_dim, 2 * channels),
|
||||
)
|
||||
nn.init.zeros_(self.mlp[-1].weight)
|
||||
nn.init.zeros_(self.mlp[-1].bias)
|
||||
|
||||
def forward(self, x: torch.Tensor, text_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [B, C, H, W]
|
||||
text_emb: [B, D_txt] or None for no-op (identity)
|
||||
Returns:
|
||||
[B, C, H, W]
|
||||
"""
|
||||
if text_emb is None:
|
||||
return x
|
||||
B = x.shape[0]
|
||||
gamma_beta = self.mlp(text_emb.to(x.dtype)) # [B, 2C]
|
||||
gamma, beta = gamma_beta.chunk(2, dim=-1)
|
||||
gamma = gamma.view(B, self.channels, 1, 1) + 1.0
|
||||
beta = beta.view(B, self.channels, 1, 1)
|
||||
return gamma * x + beta
|
||||
|
||||
|
||||
# ============================================================
|
||||
# RoPE 2D
|
||||
# ============================================================
|
||||
|
||||
class RoPE2D(nn.Module):
|
||||
"""2D Rotary Position Embedding.
|
||||
|
||||
Splits head_dim into two halves: first half gets x-position encoding,
|
||||
second half gets y-position encoding. For each half, applies standard
|
||||
1D RoPE rotation.
|
||||
|
||||
Reference: RoFormer (B49) adapted for 2D.
|
||||
"""
|
||||
|
||||
def __init__(self, head_dim: int, max_resolution: int = 64, base: float = 10000.0) -> None:
|
||||
super().__init__()
|
||||
assert head_dim % 2 == 0, "head_dim must be even for RoPE"
|
||||
assert head_dim % 4 == 0, "head_dim must be divisible by 4 for 2D RoPE"
|
||||
self.head_dim = head_dim
|
||||
self.half_dim = head_dim // 2 # dedicated to each axis
|
||||
self.max_resolution = max_resolution
|
||||
|
||||
# Frequencies for each axis (half_dim per axis, sin+cos pairs)
|
||||
freqs = 1.0 / (base ** (torch.arange(0, self.half_dim, 2).float() / self.half_dim))
|
||||
self.register_buffer("freqs", freqs, persistent=False)
|
||||
|
||||
def _make_embeds(self, H: int, W: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Produce cos/sin embeddings for HW tokens in raster order."""
|
||||
y_pos = torch.arange(H, device=device, dtype=torch.float32)
|
||||
x_pos = torch.arange(W, device=device, dtype=torch.float32)
|
||||
|
||||
freqs_y = torch.einsum("i,j->ij", y_pos, self.freqs) # [H, half_dim/2]
|
||||
freqs_x = torch.einsum("i,j->ij", x_pos, self.freqs) # [W, half_dim/2]
|
||||
|
||||
# Expand to full grid: [H, W, half_dim/2] each
|
||||
freqs_y = freqs_y.unsqueeze(1).expand(-1, W, -1) # [H, W, half_dim/2]
|
||||
freqs_x = freqs_x.unsqueeze(0).expand(H, -1, -1) # [H, W, half_dim/2]
|
||||
|
||||
# Concatenate: x-axis freqs into first half, y-axis into second half
|
||||
# Each half is [H, W, half_dim/2]; we pair-up for complex rotation
|
||||
freqs_combined_x = torch.cat([freqs_x, freqs_x], dim=-1) # [H, W, half_dim]
|
||||
freqs_combined_y = torch.cat([freqs_y, freqs_y], dim=-1) # [H, W, half_dim]
|
||||
|
||||
freqs_full = torch.cat([freqs_combined_x, freqs_combined_y], dim=-1) # [H, W, head_dim]
|
||||
cos = freqs_full.cos().reshape(H * W, -1)
|
||||
sin = freqs_full.sin().reshape(H * W, -1)
|
||||
return cos, sin
|
||||
|
||||
@staticmethod
|
||||
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Rotate: (x1, x2) -> (-x2, x1)."""
|
||||
x1, x2 = x.chunk(2, dim=-1)
|
||||
return torch.cat([-x2, x1], dim=-1)
|
||||
|
||||
def forward(self, q: torch.Tensor, k: torch.Tensor, H: int, W: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
q, k: [B, heads, HW, head_dim]
|
||||
H, W: spatial dims for position computation
|
||||
Returns:
|
||||
q_rot, k_rot with positional encoding applied
|
||||
"""
|
||||
cos, sin = self._make_embeds(H, W, q.device)
|
||||
cos = cos.to(q.dtype).unsqueeze(0).unsqueeze(0) # [1, 1, HW, head_dim]
|
||||
sin = sin.to(q.dtype).unsqueeze(0).unsqueeze(0)
|
||||
q_rot = (q * cos) + (self._rotate_half(q) * sin)
|
||||
k_rot = (k * cos) + (self._rotate_half(k) * sin)
|
||||
return q_rot, k_rot
|
||||
|
||||
|
||||
# ============================================================
|
||||
# SE: Squeeze-Excite
|
||||
# ============================================================
|
||||
|
||||
class SqueezeExcite(nn.Module):
|
||||
"""Standard Squeeze-Excite channel attention."""
|
||||
|
||||
def __init__(self, channels: int, reduction: int = 16) -> None:
|
||||
super().__init__()
|
||||
hidden = max(1, channels // reduction)
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Conv2d(channels, hidden, 1),
|
||||
nn.SiLU(inplace=True),
|
||||
nn.Conv2d(hidden, channels, 1),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
s = self.pool(x)
|
||||
s = self.fc(s)
|
||||
return x * s
|
||||
|
||||
|
||||
# ============================================================
|
||||
# LayerNorm2d: LN over channels for (B, C, H, W) layout
|
||||
# ============================================================
|
||||
|
||||
class LayerNorm2d(nn.Module):
|
||||
"""LayerNorm over C dimension for 4D tensors."""
|
||||
|
||||
def __init__(self, channels: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.norm = nn.LayerNorm(channels, eps=eps)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# [B, C, H, W] → [B, H, W, C] → LN → [B, C, H, W]
|
||||
x = x.permute(0, 2, 3, 1)
|
||||
x = self.norm(x)
|
||||
x = x.permute(0, 3, 1, 2).contiguous()
|
||||
return x
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Smoke test
|
||||
torch.manual_seed(0)
|
||||
|
||||
# GGeM test
|
||||
g = GGeM(64)
|
||||
x = torch.randn(2, 64, 8, 8)
|
||||
out = g(x)
|
||||
print(f"GGeM out: {out.shape}") # [2, 64]
|
||||
|
||||
# CHP test — verify rotation invariance
|
||||
chp = CircularHarmonicPool(32, rings=8, angles=16, harmonics=4)
|
||||
x = torch.randn(1, 32, 16, 16)
|
||||
out1 = chp(x)
|
||||
# Rotate x by 90° and verify invariance (approximately)
|
||||
x_rot = torch.rot90(x, k=1, dims=(-2, -1))
|
||||
out2 = chp(x_rot)
|
||||
diff = (out1 - out2).abs().max().item()
|
||||
print(f"CHP rotation-invariance max diff: {diff:.4e} (should be small)")
|
||||
print(f"CHP out shape: {out1.shape}") # [1, 128]
|
||||
|
||||
# FiLM test
|
||||
film = AltitudeFiLM(64)
|
||||
x = torch.randn(2, 64, 8, 8)
|
||||
altitudes = torch.tensor([150.0, 300.0])
|
||||
out = film(x, altitudes)
|
||||
print(f"FiLM out: {out.shape}") # [2, 64, 8, 8]
|
||||
|
||||
# RoPE2D test
|
||||
rope = RoPE2D(32)
|
||||
q = torch.randn(2, 4, 64, 32) # [B, heads, HW, head_dim]
|
||||
k = torch.randn(2, 4, 64, 32)
|
||||
q_r, k_r = rope(q, k, 8, 8)
|
||||
print(f"RoPE out: q={q_r.shape}, k={k_r.shape}")
|
||||
589
src/models/sofia_v71/model.py
Normal file
589
src/models/sofia_v71/model.py
Normal file
@@ -0,0 +1,589 @@
|
||||
"""SOFIA v7.1 full model.
|
||||
|
||||
Architecture:
|
||||
stem → stage1 (StripDCN-lite) → stage2 (StripMixConv)
|
||||
→ stage3 (MambaVision MV5) → stage4 (MambaVision MV1)
|
||||
→ 1×1 neck projection → {Sat-Head, UAV-Head} + (training) Ring aux
|
||||
|
||||
Features:
|
||||
- Siamese backbone with optional weight-sharing stages 1-2
|
||||
- Asymmetric Sat/UAV heads (CVGL-Aware Head v7.1-α)
|
||||
- Training-time Ring LPN aux for partial matching robustness
|
||||
- KD taps on F2, F3, F4 for future teacher-student distillation
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .config import SOFIAConfig
|
||||
from .blocks import (
|
||||
Downsample,
|
||||
EVSSBridge,
|
||||
MambaVisionBlock,
|
||||
StripDCNLiteBlock,
|
||||
StripMixConvBlock,
|
||||
)
|
||||
from .layers import (
|
||||
AltitudeFiLM,
|
||||
CircularHarmonicPool,
|
||||
GGeM,
|
||||
LayerNorm2d,
|
||||
TextFiLM,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Stem
|
||||
# ============================================================
|
||||
|
||||
class Stem(nn.Module):
|
||||
"""Dual-conv FastViT-style stem: 3 → mid (s=2) → out (s=1).
|
||||
|
||||
Downsampling ×2 total (input 256 → 128).
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels: int, mid_channels: int, out_channels: int) -> None:
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(in_channels, mid_channels, 3, stride=2, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(mid_channels)
|
||||
self.act1 = nn.SiLU(inplace=True)
|
||||
self.conv2 = nn.Conv2d(mid_channels, out_channels, 3, stride=1, padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(out_channels)
|
||||
self.act2 = nn.SiLU(inplace=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.act1(self.bn1(self.conv1(x)))
|
||||
x = self.act2(self.bn2(self.conv2(x)))
|
||||
return x
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Backbone (shared s1-2 optional, separate s3-4)
|
||||
# ============================================================
|
||||
|
||||
class Backbone(nn.Module):
|
||||
"""SOFIA v7.1 backbone: 4 stages + stem.
|
||||
|
||||
Stage 1: StripDCN-lite blocks
|
||||
Stage 2: StripMixConv blocks
|
||||
Stage 3: MambaVision MV5 blocks (Mamba ∥ MHSA ∥ Strip)
|
||||
Stage 4: MambaVision MV1 blocks (Mamba ∥ MHSA)
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: SOFIAConfig) -> None:
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
dims = cfg.embed_dims
|
||||
depths = cfg.depths
|
||||
|
||||
# Stem: input → stem_out
|
||||
self.stem = Stem(cfg.in_channels, cfg.stem_mid, cfg.stem_out)
|
||||
|
||||
# Stage 1: stem_out → dims[0], no downsample at block level
|
||||
# (stem already did 4× total; stage 1 operates at 64×64 for 256 input)
|
||||
self.ds1 = Downsample(cfg.stem_out, dims[0])
|
||||
self.stage1 = nn.Sequential(*[
|
||||
StripDCNLiteBlock(
|
||||
in_channels=dims[0],
|
||||
out_channels=dims[0],
|
||||
expand=cfg.mbconv_expand,
|
||||
kernel=cfg.strip_kernel_s1,
|
||||
stride=1,
|
||||
se_ratio=cfg.se_ratio,
|
||||
use_dcn=cfg.use_dcn_strip,
|
||||
)
|
||||
for _ in range(depths[0])
|
||||
])
|
||||
|
||||
# Stage 2
|
||||
self.ds2 = Downsample(dims[0], dims[1])
|
||||
self.stage2 = nn.Sequential(*[
|
||||
StripMixConvBlock(
|
||||
in_channels=dims[1],
|
||||
out_channels=dims[1],
|
||||
expand=cfg.mbconv_expand,
|
||||
strip_kernel=cfg.strip_kernel_s2,
|
||||
mix_kernels=cfg.mix_kernels,
|
||||
stride=1,
|
||||
se_ratio=cfg.se_ratio,
|
||||
)
|
||||
for _ in range(depths[1])
|
||||
])
|
||||
|
||||
# Stage 3 — optional EVSS bridge after downsample (B6-inspired refinement
|
||||
# to smooth DCN→MambaVision semantic gap)
|
||||
self.ds3 = Downsample(dims[1], dims[2])
|
||||
if cfg.use_evss_bridge and "pre_stage3" in cfg.evss_bridge_locations:
|
||||
self.bridge3 = EVSSBridge(
|
||||
channels=dims[2],
|
||||
mamba_d_state=cfg.mamba_d_state,
|
||||
mamba_dt_rank=cfg.mamba_dt_rank,
|
||||
mamba_variant="mamba1", # lightweight refinement, not primary mixer
|
||||
mamba_backend=cfg.mamba_backend,
|
||||
se_reduction=cfg.se_ratio,
|
||||
)
|
||||
else:
|
||||
self.bridge3 = None
|
||||
self.stage3 = nn.Sequential(*[
|
||||
MambaVisionBlock(
|
||||
channels=dims[2],
|
||||
num_heads=cfg.num_heads_s3,
|
||||
d_state=cfg.mamba_d_state,
|
||||
dt_rank=cfg.mamba_dt_rank,
|
||||
use_strip_branch=cfg.use_strip_branch_s3,
|
||||
ffn_expand=cfg.ffn_expand,
|
||||
strip_kernel=cfg.strip_kernel_s1,
|
||||
mamba_backend=cfg.mamba_backend,
|
||||
mamba_variant=cfg.mamba_variant,
|
||||
mamba_extra_kwargs=cfg.mamba_extra_kwargs,
|
||||
)
|
||||
for _ in range(depths[2])
|
||||
])
|
||||
|
||||
# Stage 4 — optional EVSS bridge after downsample
|
||||
self.ds4 = Downsample(dims[2], dims[3])
|
||||
if cfg.use_evss_bridge and "pre_stage4" in cfg.evss_bridge_locations:
|
||||
self.bridge4 = EVSSBridge(
|
||||
channels=dims[3],
|
||||
mamba_d_state=cfg.mamba_d_state,
|
||||
mamba_dt_rank=cfg.mamba_dt_rank,
|
||||
mamba_variant="mamba1",
|
||||
mamba_backend=cfg.mamba_backend,
|
||||
se_reduction=cfg.se_ratio,
|
||||
)
|
||||
else:
|
||||
self.bridge4 = None
|
||||
self.stage4 = nn.Sequential(*[
|
||||
MambaVisionBlock(
|
||||
channels=dims[3],
|
||||
num_heads=cfg.num_heads_s4,
|
||||
d_state=cfg.mamba_d_state,
|
||||
dt_rank=cfg.mamba_dt_rank,
|
||||
use_strip_branch=cfg.use_strip_branch_s4,
|
||||
ffn_expand=cfg.ffn_expand,
|
||||
strip_kernel=cfg.strip_kernel_s1,
|
||||
mamba_backend=cfg.mamba_backend,
|
||||
mamba_variant=cfg.mamba_variant,
|
||||
mamba_extra_kwargs=cfg.mamba_extra_kwargs,
|
||||
)
|
||||
for _ in range(depths[3])
|
||||
])
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
|
||||
"""Returns dict with F1-F4 taps for KD/multi-scale use."""
|
||||
s0 = self.stem(x) # stem_out, H/2
|
||||
f1 = self.stage1(self.ds1(s0)) # dims[0], H/4
|
||||
f2 = self.stage2(self.ds2(f1)) # dims[1], H/8
|
||||
|
||||
# Stage 3 with optional EVSS bridge.
|
||||
ds3_out = self.ds3(f2)
|
||||
if self.bridge3 is not None:
|
||||
ds3_out = self.bridge3(ds3_out)
|
||||
f3 = self.stage3(ds3_out) # dims[2], H/16
|
||||
|
||||
# Stage 4 with optional EVSS bridge.
|
||||
ds4_out = self.ds4(f3)
|
||||
if self.bridge4 is not None:
|
||||
ds4_out = self.bridge4(ds4_out)
|
||||
f4 = self.stage4(ds4_out) # dims[3], H/32 (= 8×8 for 256 input)
|
||||
|
||||
return {"s0": s0, "f1": f1, "f2": f2, "f3": f3, "f4": f4}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Neck: ultra-lite 1×1 projection
|
||||
# ============================================================
|
||||
|
||||
class UltraLiteNeck(nn.Module):
|
||||
"""1×1 projection + BN + activation."""
|
||||
|
||||
def __init__(self, in_channels: int, out_channels: int) -> None:
|
||||
super().__init__()
|
||||
self.proj = nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, 1, bias=False),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.SiLU(inplace=True),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.proj(x)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Heads
|
||||
# ============================================================
|
||||
|
||||
class SatHead(nn.Module):
|
||||
"""Satellite view head: [TextFiLM] + GGeM + BN + Linear [+ L2].
|
||||
|
||||
`text_emb` is optional; when None or `text_film` is disabled, behaves
|
||||
identically to the original head. `return_normalized=False` returns
|
||||
pre-L2 descriptors for late gated fusion in a wrapper.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
d_descriptor: int,
|
||||
return_normalized: bool = True,
|
||||
use_text_film: bool = False,
|
||||
text_film_dim: int = 1024,
|
||||
text_film_hidden: int = 256,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.return_normalized = return_normalized
|
||||
self.use_text_film = use_text_film
|
||||
if use_text_film:
|
||||
self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
|
||||
self.ggem = GGeM(channels)
|
||||
self.bn = nn.BatchNorm1d(channels, affine=False)
|
||||
self.proj = nn.Linear(channels, d_descriptor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
text_emb: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""x: [B, C, H, W], text_emb: [B, D_txt] or None -> [B, d]"""
|
||||
if self.use_text_film:
|
||||
x = self.text_film(x, text_emb)
|
||||
g = self.ggem(x)
|
||||
g = self.bn(g)
|
||||
g = self.proj(g)
|
||||
if self.return_normalized:
|
||||
g = F.normalize(g, p=2, dim=-1)
|
||||
return g
|
||||
|
||||
|
||||
class UAVHead(nn.Module):
|
||||
"""UAV view head: FiLM(altitude) [+ TextFiLM] + CHP + BN + Linear [+ L2].
|
||||
|
||||
Formally SO(2)-invariant via CHP. Altitude-aware via FiLM. Optional
|
||||
text-FiLM is applied AFTER altitude-FiLM (zero-init β so it starts as
|
||||
identity). `return_normalized=False` returns pre-L2 descriptors for
|
||||
late gated fusion in a wrapper.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
d_descriptor: int,
|
||||
rings: int = 8,
|
||||
angles: int = 16,
|
||||
harmonics: int = 4,
|
||||
use_film: bool = True,
|
||||
altitude_norm: float = 500.0,
|
||||
return_normalized: bool = True,
|
||||
use_text_film: bool = False,
|
||||
text_film_dim: int = 1024,
|
||||
text_film_hidden: int = 256,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.return_normalized = return_normalized
|
||||
self.use_film = use_film
|
||||
self.use_text_film = use_text_film
|
||||
if use_film:
|
||||
self.film = AltitudeFiLM(channels, altitude_norm=altitude_norm)
|
||||
if use_text_film:
|
||||
self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
|
||||
|
||||
self.chp = CircularHarmonicPool(
|
||||
channels=channels,
|
||||
rings=rings,
|
||||
angles=angles,
|
||||
harmonics=harmonics,
|
||||
)
|
||||
chp_dim = channels * harmonics
|
||||
self.bn = nn.BatchNorm1d(chp_dim, affine=False)
|
||||
self.proj = nn.Linear(chp_dim, d_descriptor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
altitude: Optional[torch.Tensor] = None,
|
||||
text_emb: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""x: [B, C, H, W], altitude: [B] or None, text_emb: [B, D_txt] or None -> [B, d]"""
|
||||
if self.use_film:
|
||||
x = self.film(x, altitude)
|
||||
if self.use_text_film:
|
||||
x = self.text_film(x, text_emb)
|
||||
g = self.chp(x)
|
||||
g = self.bn(g)
|
||||
g = self.proj(g)
|
||||
if self.return_normalized:
|
||||
g = F.normalize(g, p=2, dim=-1)
|
||||
return g
|
||||
|
||||
|
||||
class RingAuxHead(nn.Module):
|
||||
"""LPN Square-Ring pool + per-ring Linear + L2 (training-only auxiliary).
|
||||
|
||||
Used for partial-matching robustness. Drop at inference.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
rings: int = 4,
|
||||
d_per_ring: int = 128,
|
||||
feature_size: int = 8,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.rings = rings
|
||||
self.feature_size = feature_size
|
||||
self.d_per_ring = d_per_ring
|
||||
|
||||
# Per-ring GGeM
|
||||
self.ggems = nn.ModuleList([GGeM(channels) for _ in range(rings)])
|
||||
self.projs = nn.ModuleList([
|
||||
nn.Linear(channels, d_per_ring) for _ in range(rings)
|
||||
])
|
||||
|
||||
# Precompute ring masks
|
||||
masks = self._make_ring_masks(rings, feature_size) # [R, H, W]
|
||||
self.register_buffer("ring_masks", masks, persistent=False)
|
||||
|
||||
@staticmethod
|
||||
def _make_ring_masks(R: int, S: int) -> torch.Tensor:
|
||||
"""Concentric square rings on SxS feature map."""
|
||||
masks = torch.zeros(R, S, S)
|
||||
center = (S - 1) / 2.0
|
||||
# Define ring by Chebyshev distance thresholds
|
||||
for r in range(R):
|
||||
r_min = r * (S / (2 * R))
|
||||
r_max = (r + 1) * (S / (2 * R))
|
||||
for i in range(S):
|
||||
for j in range(S):
|
||||
dist = max(abs(i - center), abs(j - center))
|
||||
if r_min <= dist < r_max or (r == R - 1 and dist <= r_max):
|
||||
masks[r, i, j] = 1.0
|
||||
# Normalize each ring to sum to 1 for pooling stability (not required but cleaner)
|
||||
masks = masks / masks.sum(dim=(1, 2), keepdim=True).clamp(min=1.0)
|
||||
return masks
|
||||
|
||||
def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
x: [B, C, H, W]
|
||||
Returns:
|
||||
List of R tensors each [B, d_per_ring], L2-normalized
|
||||
"""
|
||||
B, C, H, W = x.shape
|
||||
outs = []
|
||||
for r in range(self.rings):
|
||||
mask = self.ring_masks[r].to(x.dtype) # [H, W]
|
||||
# Apply mask
|
||||
x_masked = x * mask.unsqueeze(0).unsqueeze(0) # [B, C, H, W]
|
||||
g = self.ggems[r](x_masked) # [B, C]
|
||||
g = self.projs[r](g) # [B, d_per_ring]
|
||||
g = F.normalize(g, p=2, dim=-1)
|
||||
outs.append(g)
|
||||
return outs
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Full SOFIA v7.1 model
|
||||
# ============================================================
|
||||
|
||||
class SOFIAv71(nn.Module):
|
||||
"""Full SOFIA v7.1 model.
|
||||
|
||||
Forward signature:
|
||||
forward(sat: [B,3,H,W], uav: [B,3,H,W], altitude: [B] = None,
|
||||
return_features: bool = False) -> dict
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: SOFIAConfig) -> None:
|
||||
super().__init__()
|
||||
cfg.validate()
|
||||
self.cfg = cfg
|
||||
|
||||
# Backbones (shared or separate depending on config)
|
||||
self.backbone_shared = Backbone(cfg)
|
||||
if cfg.share_stages_1_2:
|
||||
# Single backbone, use one for both
|
||||
self.backbone_sat = None
|
||||
self.backbone_uav = None
|
||||
else:
|
||||
# Fully separate backbones (rare case)
|
||||
self.backbone_sat = Backbone(cfg)
|
||||
self.backbone_uav = Backbone(cfg)
|
||||
# Clear shared
|
||||
del self.backbone_shared
|
||||
self.backbone_shared = None
|
||||
|
||||
# Neck (shared)
|
||||
self.neck = UltraLiteNeck(cfg.embed_dims[-1], cfg.neck_channels)
|
||||
|
||||
# Heads
|
||||
if cfg.use_asymmetric_heads:
|
||||
self.sat_head = SatHead(
|
||||
cfg.neck_channels, cfg.d_descriptor,
|
||||
return_normalized=cfg.return_normalized,
|
||||
use_text_film=cfg.use_text_film_sat,
|
||||
text_film_dim=cfg.text_film_dim,
|
||||
text_film_hidden=cfg.text_film_hidden,
|
||||
)
|
||||
self.uav_head = UAVHead(
|
||||
channels=cfg.neck_channels,
|
||||
d_descriptor=cfg.d_descriptor,
|
||||
rings=cfg.chp_rings,
|
||||
angles=cfg.chp_angles,
|
||||
harmonics=cfg.chp_harmonics,
|
||||
use_film=cfg.use_film_altitude,
|
||||
altitude_norm=cfg.altitude_norm,
|
||||
return_normalized=cfg.return_normalized,
|
||||
use_text_film=cfg.use_text_film_uav,
|
||||
text_film_dim=cfg.text_film_dim,
|
||||
text_film_hidden=cfg.text_film_hidden,
|
||||
)
|
||||
else:
|
||||
# Symmetric: use SatHead for both
|
||||
self.sat_head = SatHead(
|
||||
cfg.neck_channels, cfg.d_descriptor,
|
||||
return_normalized=cfg.return_normalized,
|
||||
use_text_film=cfg.use_text_film_sat,
|
||||
text_film_dim=cfg.text_film_dim,
|
||||
text_film_hidden=cfg.text_film_hidden,
|
||||
)
|
||||
self.uav_head = self.sat_head
|
||||
|
||||
# Ring aux (training-only)
|
||||
# Resolution chain: input/2 (stem) → /2 (ds1) → /2 (ds2) → /2 (ds3) → /2 (ds4)
|
||||
# Total ×32 downsampling → input 256 gives 8×8 final feature map
|
||||
if cfg.use_ring_aux:
|
||||
feat_size = cfg.input_size // 32
|
||||
self.ring_head = RingAuxHead(
|
||||
channels=cfg.neck_channels,
|
||||
rings=cfg.ring_count,
|
||||
d_per_ring=cfg.d_descriptor // cfg.ring_count,
|
||||
feature_size=feat_size,
|
||||
)
|
||||
else:
|
||||
self.ring_head = None
|
||||
|
||||
def _extract_features(self, x: torch.Tensor, view: str) -> Dict[str, torch.Tensor]:
|
||||
"""Run backbone for sat or uav view."""
|
||||
if self.cfg.share_stages_1_2:
|
||||
return self.backbone_shared(x)
|
||||
else:
|
||||
bb = self.backbone_sat if view == "sat" else self.backbone_uav
|
||||
return bb(x)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sat: Optional[torch.Tensor] = None,
|
||||
uav: Optional[torch.Tensor] = None,
|
||||
altitude: Optional[torch.Tensor] = None,
|
||||
text_emb_sat: Optional[torch.Tensor] = None,
|
||||
text_emb_uav: Optional[torch.Tensor] = None,
|
||||
return_features: bool = False,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
sat: [B, 3, H, W] satellite image (or None to skip)
|
||||
uav: [B, 3, H, W] UAV image (or None to skip)
|
||||
altitude: [B] altitude in meters for UAV (or None = neutral)
|
||||
text_emb_sat: [B, D_txt] caption embedding for SatHead text-FiLM (or None)
|
||||
text_emb_uav: [B, D_txt] caption embedding for UAVHead text-FiLM (or None)
|
||||
return_features: if True, also return backbone features F2-F4 for KD
|
||||
|
||||
Returns:
|
||||
dict with g_sat, g_uav (global descriptors), optional features,
|
||||
and optional rings (training only)
|
||||
"""
|
||||
result: Dict[str, torch.Tensor] = {}
|
||||
|
||||
# SAT path
|
||||
if sat is not None:
|
||||
feats_sat = self._extract_features(sat, view="sat")
|
||||
f4_sat = feats_sat["f4"]
|
||||
neck_sat = self.neck(f4_sat)
|
||||
g_sat = self.sat_head(neck_sat, text_emb=text_emb_sat)
|
||||
result["g_sat"] = g_sat
|
||||
if return_features:
|
||||
result["features_sat"] = feats_sat
|
||||
result["neck_sat"] = neck_sat
|
||||
if self.training and self.ring_head is not None:
|
||||
result["rings_sat"] = self.ring_head(neck_sat)
|
||||
|
||||
# UAV path
|
||||
if uav is not None:
|
||||
feats_uav = self._extract_features(uav, view="uav")
|
||||
f4_uav = feats_uav["f4"]
|
||||
neck_uav = self.neck(f4_uav)
|
||||
if self.cfg.use_asymmetric_heads:
|
||||
g_uav = self.uav_head(neck_uav, altitude=altitude, text_emb=text_emb_uav)
|
||||
else:
|
||||
g_uav = self.uav_head(neck_uav, text_emb=text_emb_uav)
|
||||
result["g_uav"] = g_uav
|
||||
if return_features:
|
||||
result["features_uav"] = feats_uav
|
||||
result["neck_uav"] = neck_uav
|
||||
if self.training and self.ring_head is not None:
|
||||
result["rings_uav"] = self.ring_head(neck_uav)
|
||||
|
||||
return result
|
||||
|
||||
def count_parameters(self, only_trainable: bool = True) -> Dict[str, int]:
|
||||
"""Count parameters per module."""
|
||||
counts = {}
|
||||
total = 0
|
||||
for name, module in self.named_children():
|
||||
if module is None:
|
||||
continue
|
||||
n = sum(p.numel() for p in module.parameters() if not only_trainable or p.requires_grad)
|
||||
counts[name] = n
|
||||
total += n
|
||||
counts["_total"] = total
|
||||
return counts
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Build helper
|
||||
# ============================================================
|
||||
|
||||
def build_sofia(preset: str = "M") -> SOFIAv71:
|
||||
"""Build SOFIA v7.1 model from preset name."""
|
||||
from .config import sofia_m_config, sofia_l_config, sofia_tiny_config
|
||||
preset_map = {
|
||||
"M": sofia_m_config,
|
||||
"L": sofia_l_config,
|
||||
"Tiny": sofia_tiny_config,
|
||||
}
|
||||
if preset not in preset_map:
|
||||
raise ValueError(f"Unknown preset '{preset}'. Available: {list(preset_map)}")
|
||||
cfg = preset_map[preset]()
|
||||
return SOFIAv71(cfg)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Smoke test with Tiny preset (fast)
|
||||
print("Building SOFIA-Tiny for smoke test...")
|
||||
model = build_sofia("Tiny")
|
||||
model.eval()
|
||||
|
||||
counts = model.count_parameters()
|
||||
total_m = counts["_total"] / 1e6
|
||||
print(f"Total params: {total_m:.2f} M")
|
||||
print("Per-module:")
|
||||
for k, v in counts.items():
|
||||
print(f" {k}: {v / 1e6:.3f} M")
|
||||
|
||||
# Dry forward
|
||||
sat = torch.randn(1, 3, 256, 256)
|
||||
uav = torch.randn(1, 3, 256, 256)
|
||||
alt = torch.tensor([150.0])
|
||||
with torch.no_grad():
|
||||
out = model(sat=sat, uav=uav, altitude=alt)
|
||||
for k, v in out.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
print(f" out[{k}]: {tuple(v.shape)}")
|
||||
else:
|
||||
print(f" out[{k}]: {type(v).__name__}")
|
||||
565
src/models/sofia_v71/quant.py
Normal file
565
src/models/sofia_v71/quant.py
Normal file
@@ -0,0 +1,565 @@
|
||||
"""SOFIA v7.1 quantization utilities.
|
||||
|
||||
Two reusable building blocks for INT8 deployment:
|
||||
|
||||
1. **OffsetClampSTE** (DCN-M2): hard clamp DCN offsets to physical range
|
||||
``[-k, +k]`` with straight-through gradient. Bounded distribution =
|
||||
stable percentile clipping during PTQ + cleaner offset semantics.
|
||||
|
||||
2. **KScaledFakeQuant** (R5-style): multi-bin fake quantization. Bins are
|
||||
chosen by magnitude percentiles; each bin has its own scale. Adapts
|
||||
quantization resolution to long-tail distributions (Mamba Δ, output y).
|
||||
|
||||
Plus a drop-in:
|
||||
|
||||
3. **KScaledMamba2Block**: ``Mamba2Block`` subclass with k-scaled fake-quant
|
||||
nodes on ``x_main``, ``delta``, and per-token ``y`` outputs. Implements
|
||||
R5 reparam principle (state ``h_t`` kept in model dtype; only observable
|
||||
tensors quantized).
|
||||
|
||||
All components are PTQ/QAT-compatible (STE backward, calibration mode).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Optional, Sequence, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .blocks import Mamba2Block
|
||||
|
||||
|
||||
# ============================================================
|
||||
# DCN-M2: Offset clamp with STE
|
||||
# ============================================================
|
||||
|
||||
class _OffsetClampSTE(torch.autograd.Function):
|
||||
"""Hard clamp [low, high] in forward, identity gradient in backward."""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x: torch.Tensor, low: float, high: float) -> torch.Tensor:
|
||||
return torch.clamp(x, low, high)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad: torch.Tensor):
|
||||
# STE: pass gradient unmodified
|
||||
return grad, None, None
|
||||
|
||||
|
||||
def offset_clamp_ste(offsets: torch.Tensor, kernel_size: int) -> torch.Tensor:
|
||||
"""Functional DCN-M2: clamp offsets to ``[-k, +k]`` with STE backward.
|
||||
|
||||
Args:
|
||||
offsets: arbitrary-shape tensor of DCN offsets.
|
||||
kernel_size: physical receptive-field bound k. Output is clamped to
|
||||
absolute values <= k. Standard choice equals the DW kernel size.
|
||||
|
||||
Returns:
|
||||
Clamped tensor (same shape as ``offsets``). Backward propagates
|
||||
gradient as if no clamp was applied, so training can still adjust
|
||||
weights even when an offset hits the boundary.
|
||||
"""
|
||||
k = float(kernel_size)
|
||||
return _OffsetClampSTE.apply(offsets, -k, k)
|
||||
|
||||
|
||||
class OffsetClampSTE(nn.Module):
|
||||
"""Module wrapper for :func:`offset_clamp_ste`.
|
||||
|
||||
Use as a drop-in layer between ``offset_predictor`` and ``deform_conv2d``::
|
||||
|
||||
self.clamp = OffsetClampSTE(kernel_size=7)
|
||||
offsets = self.clamp(self.offset_predictor(x))
|
||||
"""
|
||||
|
||||
def __init__(self, kernel_size: int) -> None:
|
||||
super().__init__()
|
||||
self.kernel_size = int(kernel_size)
|
||||
|
||||
def forward(self, offsets: torch.Tensor) -> torch.Tensor:
|
||||
return offset_clamp_ste(offsets, self.kernel_size)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"kernel_size={self.kernel_size}, range=[-{self.kernel_size}, +{self.kernel_size}]"
|
||||
|
||||
|
||||
# ============================================================
|
||||
# K-Scaled Fake Quantization (R5 style, multi-bin)
|
||||
# ============================================================
|
||||
|
||||
class _STEFakeQuant(torch.autograd.Function):
|
||||
"""Fake-quantize: ``round(x / s) * s`` clamped to [qmin, qmax], STE backward.
|
||||
|
||||
`scale` is a tensor broadcastable with `x` (per-tensor scalar or
|
||||
per-element via gathered per-bin scales).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
qmin: int,
|
||||
qmax: int,
|
||||
) -> torch.Tensor:
|
||||
x_int = torch.round(x / scale).clamp(qmin, qmax)
|
||||
return x_int * scale
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output: torch.Tensor):
|
||||
return grad_output, None, None, None
|
||||
|
||||
|
||||
class KScaledFakeQuant(nn.Module):
|
||||
"""k-scaled (multi-bin magnitude-bucketed) fake quantization.
|
||||
|
||||
Algorithm:
|
||||
1. ``g(x) = bin index in [0..k-1] based on |x|`` using thresholds
|
||||
``t_0 < t_1 < ... < t_{k-2}``.
|
||||
2. ``s(x) = scales[g(x)]`` per-element scale.
|
||||
3. ``x_q = round(x / s) * s`` clamped to ``[qmin, qmax]``.
|
||||
4. STE backward.
|
||||
|
||||
Setting ``num_bins=1`` reduces to standard single-scale fake quant.
|
||||
|
||||
Calibration workflow::
|
||||
|
||||
fq = KScaledFakeQuant(num_bins=3)
|
||||
fq.start_calibration()
|
||||
for batch in cal_loader:
|
||||
_ = model(batch) # collect stats
|
||||
fq.finalize_calibration(percentile=99.9)
|
||||
# fq is now active (calibrating=False) and applies fake-quant in forward
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_bins: int = 3,
|
||||
qmin: int = -128,
|
||||
qmax: int = 127,
|
||||
max_calibration_samples: int = 200_000,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if num_bins < 1:
|
||||
raise ValueError(f"num_bins must be >= 1, got {num_bins}")
|
||||
if qmax <= qmin:
|
||||
raise ValueError(f"qmax ({qmax}) must be > qmin ({qmin})")
|
||||
|
||||
self.num_bins = num_bins
|
||||
self.qmin = qmin
|
||||
self.qmax = qmax
|
||||
self.max_calibration_samples = max_calibration_samples
|
||||
|
||||
# Thresholds: (num_bins - 1,) — bin boundaries on |x|. For num_bins=1, empty.
|
||||
n_thr = max(0, num_bins - 1)
|
||||
self.register_buffer("thresholds", torch.zeros(n_thr))
|
||||
|
||||
# Scales: (num_bins,) — one scale per bin. Default 1.0 = no quant effect.
|
||||
self.register_buffer("scales", torch.ones(num_bins))
|
||||
|
||||
# Per-instance state (not buffers — runtime-only)
|
||||
self.calibrating: bool = False
|
||||
self.enabled: bool = True
|
||||
self._calib_buffer: List[torch.Tensor] = []
|
||||
self._n_collected: int = 0
|
||||
|
||||
# -------------------- Calibration API --------------------
|
||||
|
||||
def start_calibration(self) -> None:
|
||||
"""Begin collecting input distribution stats."""
|
||||
self.calibrating = True
|
||||
self._calib_buffer = []
|
||||
self._n_collected = 0
|
||||
|
||||
def stop_calibration(self) -> None:
|
||||
"""End calibration without finalizing scales (e.g. abort)."""
|
||||
self.calibrating = False
|
||||
self._calib_buffer = []
|
||||
self._n_collected = 0
|
||||
|
||||
def finalize_calibration(self, percentile: float = 99.9) -> None:
|
||||
"""Compute thresholds and per-bin scales from collected stats.
|
||||
|
||||
Args:
|
||||
percentile: per-bin tail percentile for scale computation
|
||||
(typical 99.9). Robust to extreme outliers.
|
||||
"""
|
||||
if not self._calib_buffer:
|
||||
self.calibrating = False
|
||||
return
|
||||
|
||||
all_data = torch.cat(self._calib_buffer)
|
||||
abs_data = all_data.abs()
|
||||
|
||||
if self.num_bins == 1:
|
||||
# Single scale: percentile of |x|
|
||||
tail = torch.quantile(abs_data, percentile / 100.0).item()
|
||||
self.scales[0] = max(tail / float(self.qmax), 1e-8)
|
||||
else:
|
||||
# Threshold positions: equal-mass quantile splits
|
||||
qpts = torch.linspace(0.0, 1.0, self.num_bins + 1, device=abs_data.device)[1:-1]
|
||||
thr_vals = torch.quantile(abs_data, qpts)
|
||||
self.thresholds.copy_(thr_vals.to(self.thresholds.device))
|
||||
|
||||
# Per-bin scale: per-bin tail percentile / qmax
|
||||
for g in range(self.num_bins):
|
||||
if g == 0:
|
||||
lo = 0.0
|
||||
hi = thr_vals[0].item()
|
||||
elif g == self.num_bins - 1:
|
||||
lo = thr_vals[-1].item()
|
||||
hi = float("inf")
|
||||
else:
|
||||
lo = thr_vals[g - 1].item()
|
||||
hi = thr_vals[g].item()
|
||||
|
||||
mask = (abs_data >= lo) & (abs_data < hi)
|
||||
if mask.any():
|
||||
bin_tail = torch.quantile(
|
||||
abs_data[mask], percentile / 100.0
|
||||
).item()
|
||||
self.scales[g] = max(bin_tail / float(self.qmax), 1e-8)
|
||||
else:
|
||||
self.scales[g] = 1.0
|
||||
|
||||
self._calib_buffer = []
|
||||
self._n_collected = 0
|
||||
self.calibrating = False
|
||||
|
||||
# -------------------- Forward --------------------
|
||||
|
||||
def _bin_index(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Per-element bin index (long tensor, shape == x.shape)."""
|
||||
if self.num_bins == 1:
|
||||
return torch.zeros_like(x, dtype=torch.long)
|
||||
|
||||
abs_x = x.abs()
|
||||
idx = torch.zeros_like(x, dtype=torch.long)
|
||||
for i in range(self.num_bins - 1):
|
||||
t = self.thresholds[i]
|
||||
idx = torch.where(abs_x >= t, idx + 1, idx)
|
||||
return idx
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if not self.enabled:
|
||||
return x
|
||||
|
||||
if self.calibrating:
|
||||
# Subsample to bound memory
|
||||
with torch.no_grad():
|
||||
flat = x.detach().flatten()
|
||||
budget = self.max_calibration_samples - self._n_collected
|
||||
if budget > 0:
|
||||
take = min(flat.numel(), budget)
|
||||
if take < flat.numel():
|
||||
# Random subsample for representativeness
|
||||
idx = torch.randperm(flat.numel(), device=flat.device)[:take]
|
||||
sample = flat[idx]
|
||||
else:
|
||||
sample = flat
|
||||
self._calib_buffer.append(sample.cpu())
|
||||
self._n_collected += sample.numel()
|
||||
return x
|
||||
|
||||
# Apply k-scaled fake-quant
|
||||
bin_idx = self._bin_index(x)
|
||||
scale = self.scales[bin_idx] # shape == x.shape
|
||||
return _STEFakeQuant.apply(x, scale, self.qmin, self.qmax)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return (
|
||||
f"num_bins={self.num_bins}, qmin={self.qmin}, qmax={self.qmax}, "
|
||||
f"calibrating={self.calibrating}, enabled={self.enabled}"
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# K-Scaled Mamba-2 drop-in
|
||||
# ============================================================
|
||||
|
||||
class KScaledMamba2Block(Mamba2Block):
|
||||
"""Drop-in replacement for :class:`Mamba2Block` with k-scaled fake-quant.
|
||||
|
||||
Adds ``KScaledFakeQuant`` nodes on three observable tensors inside the
|
||||
Mamba-2 forward pass:
|
||||
|
||||
- ``x_main``: scan input after conv1d + SiLU
|
||||
- ``delta``: time-step ``Δ`` after softplus (addresses MF2 long tail)
|
||||
- ``y``: per-token scan output (addresses MF4 state propagation)
|
||||
|
||||
The recurrent state ``h_t`` is **not** quantized — kept in model dtype
|
||||
inside the scan loop. This implements R5's reparam principle: only
|
||||
observable I/O is quantized; internal state stays high-precision.
|
||||
|
||||
Backend constraint: forces ``backend='torch'``. The mamba_ssm CUDA
|
||||
backend would require k-scaled support inside the custom kernel (out of
|
||||
scope here). For a fast deploy path, train+QAT in torch, then export to
|
||||
a custom TRT plugin that does the same INT8 scan.
|
||||
|
||||
Args:
|
||||
channels: model dim ``C``.
|
||||
d_state: state size ``N`` (Mamba-2 default 64).
|
||||
headdim: head dim (must divide channels).
|
||||
d_conv: local conv kernel.
|
||||
expand: inner expansion factor.
|
||||
backend: must be ``'torch'`` (or ``'auto'``, which resolves to torch).
|
||||
num_bins: ``k`` for k-scaled (typical 2-4).
|
||||
targets: which tensors to quantize. Subset of
|
||||
``('x_main', 'delta', 'y')``.
|
||||
|
||||
Example::
|
||||
|
||||
mamba = KScaledMamba2Block(channels=192, num_bins=3,
|
||||
targets=('delta', 'y'))
|
||||
mamba.start_calibration()
|
||||
for batch in cal_loader:
|
||||
_ = model_with_mamba(batch)
|
||||
mamba.finalize_calibration()
|
||||
# Now in PTQ mode — forward applies fake-quant on chosen targets.
|
||||
"""
|
||||
|
||||
SUPPORTED_TARGETS: Tuple[str, ...] = ("x_main", "delta", "y")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
d_state: int = 64,
|
||||
headdim: int = 64,
|
||||
d_conv: int = 4,
|
||||
expand: int = 2,
|
||||
backend: str = "torch",
|
||||
num_bins: int = 3,
|
||||
targets: Sequence[str] = ("x_main", "delta", "y"),
|
||||
) -> None:
|
||||
# Force torch — k-scaled requires explicit forward-pass control
|
||||
if backend == "auto":
|
||||
backend = "torch"
|
||||
if backend != "torch":
|
||||
raise ValueError(
|
||||
"KScaledMamba2Block currently supports only backend='torch'. "
|
||||
f"Got backend='{backend}'. The mamba_ssm CUDA kernel does not "
|
||||
"expose hooks for k-scaled fake-quant; integrating INT8 scan "
|
||||
"via TRT plugin is a separate deploy-time concern."
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
channels=channels,
|
||||
d_state=d_state,
|
||||
headdim=headdim,
|
||||
d_conv=d_conv,
|
||||
expand=expand,
|
||||
backend="torch",
|
||||
)
|
||||
|
||||
unknown = set(targets) - set(self.SUPPORTED_TARGETS)
|
||||
if unknown:
|
||||
raise ValueError(
|
||||
f"Unknown k-scaled targets: {unknown}. "
|
||||
f"Supported: {self.SUPPORTED_TARGETS}"
|
||||
)
|
||||
self.targets = set(targets)
|
||||
self.num_bins = num_bins
|
||||
|
||||
# Build fake-quant nodes only for active targets
|
||||
if "x_main" in self.targets:
|
||||
self.fq_x_main = KScaledFakeQuant(num_bins=num_bins)
|
||||
if "delta" in self.targets:
|
||||
self.fq_delta = KScaledFakeQuant(num_bins=num_bins)
|
||||
if "y" in self.targets:
|
||||
self.fq_y = KScaledFakeQuant(num_bins=num_bins)
|
||||
|
||||
# -------------------- Calibration helpers --------------------
|
||||
|
||||
def _all_fqs(self) -> List[KScaledFakeQuant]:
|
||||
out: List[KScaledFakeQuant] = []
|
||||
for name in self.SUPPORTED_TARGETS:
|
||||
attr = f"fq_{name}"
|
||||
mod = getattr(self, attr, None)
|
||||
if isinstance(mod, KScaledFakeQuant):
|
||||
out.append(mod)
|
||||
return out
|
||||
|
||||
def start_calibration(self) -> None:
|
||||
for fq in self._all_fqs():
|
||||
fq.start_calibration()
|
||||
|
||||
def finalize_calibration(self, percentile: float = 99.9) -> None:
|
||||
for fq in self._all_fqs():
|
||||
fq.finalize_calibration(percentile=percentile)
|
||||
|
||||
def set_quant_enabled(self, enabled: bool) -> None:
|
||||
for fq in self._all_fqs():
|
||||
fq.enabled = enabled
|
||||
|
||||
# -------------------- Forward --------------------
|
||||
|
||||
def _forward_torch(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Mamba-2 torch forward with k-scaled fake-quant on selected paths."""
|
||||
B, L, C = x.shape
|
||||
|
||||
# In projection (matches parent layout)
|
||||
z_in = self.in_proj(x)
|
||||
xz, BC, dt = torch.split(
|
||||
z_in,
|
||||
[2 * self.d_inner, 2 * self.d_state, self.nheads],
|
||||
dim=-1,
|
||||
)
|
||||
x_main, z_gate = xz.chunk(2, dim=-1)
|
||||
B_p, C_p = BC.chunk(2, dim=-1)
|
||||
|
||||
# Local conv + SiLU
|
||||
x_main_t = x_main.transpose(1, 2)
|
||||
x_main_t = self.conv1d(x_main_t)[..., :L]
|
||||
x_main = F.silu(x_main_t).transpose(1, 2)
|
||||
|
||||
# K-SCALED QUANT: x_main (scan input)
|
||||
if "x_main" in self.targets:
|
||||
x_main = self.fq_x_main(x_main)
|
||||
|
||||
# Δ
|
||||
dt = F.softplus(dt)
|
||||
if "delta" in self.targets:
|
||||
dt = self.fq_delta(dt)
|
||||
|
||||
A = -torch.exp(self.A_log) # [nheads]
|
||||
x_head = x_main.view(B, L, self.nheads, self.headdim)
|
||||
|
||||
# Sequential scan. R5 reparam: state h kept in model dtype throughout.
|
||||
h = torch.zeros(
|
||||
B, self.nheads, self.headdim, self.d_state,
|
||||
device=x.device, dtype=x.dtype,
|
||||
)
|
||||
ys: List[torch.Tensor] = []
|
||||
for t in range(L):
|
||||
dt_t = dt[:, t]
|
||||
B_t = B_p[:, t]
|
||||
C_t = C_p[:, t]
|
||||
x_t = x_head[:, t]
|
||||
|
||||
dA = torch.exp(dt_t * A.unsqueeze(0)).unsqueeze(-1).unsqueeze(-1)
|
||||
dB = (dt_t.unsqueeze(-1) * B_t.unsqueeze(1)).unsqueeze(2)
|
||||
|
||||
h = dA * h + dB * x_t.unsqueeze(-1)
|
||||
y_t = (h * C_t.unsqueeze(1).unsqueeze(1)).sum(dim=-1)
|
||||
y_t = y_t + self.D.unsqueeze(0).unsqueeze(-1) * x_t
|
||||
|
||||
# K-SCALED QUANT: per-token y (after scan, before gate)
|
||||
if "y" in self.targets:
|
||||
y_t = self.fq_y(y_t)
|
||||
|
||||
ys.append(y_t)
|
||||
|
||||
y = torch.stack(ys, dim=1)
|
||||
y = y.reshape(B, L, self.d_inner)
|
||||
y = y * F.silu(z_gate)
|
||||
y = self.out_proj(y)
|
||||
return y
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Model-wide calibration helpers
|
||||
# ============================================================
|
||||
|
||||
def start_calibration(model: nn.Module) -> None:
|
||||
"""Walk model, start calibration on every quant module."""
|
||||
seen = set()
|
||||
# First, KScaledMamba2Block instances drive their internal FQs.
|
||||
for m in model.modules():
|
||||
if isinstance(m, KScaledMamba2Block):
|
||||
m.start_calibration()
|
||||
for fq in m._all_fqs():
|
||||
seen.add(id(fq))
|
||||
# Then standalone KScaledFakeQuant not owned by a KScaledMamba2Block.
|
||||
for m in model.modules():
|
||||
if isinstance(m, KScaledFakeQuant) and id(m) not in seen:
|
||||
m.start_calibration()
|
||||
|
||||
|
||||
def finalize_calibration(model: nn.Module, percentile: float = 99.9) -> None:
|
||||
"""Walk model, finalize all quant modules with the same percentile."""
|
||||
seen = set()
|
||||
for m in model.modules():
|
||||
if isinstance(m, KScaledMamba2Block):
|
||||
m.finalize_calibration(percentile=percentile)
|
||||
for fq in m._all_fqs():
|
||||
seen.add(id(fq))
|
||||
for m in model.modules():
|
||||
if isinstance(m, KScaledFakeQuant) and id(m) not in seen:
|
||||
m.finalize_calibration(percentile=percentile)
|
||||
|
||||
|
||||
def set_quant_enabled(model: nn.Module, enabled: bool) -> None:
|
||||
"""Toggle every quant module on/off. Useful for FP-vs-INT8 comparison."""
|
||||
for m in model.modules():
|
||||
if isinstance(m, KScaledFakeQuant):
|
||||
m.enabled = enabled
|
||||
elif isinstance(m, KScaledMamba2Block):
|
||||
m.set_quant_enabled(enabled)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Smoke tests
|
||||
# ============================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
torch.manual_seed(0)
|
||||
|
||||
print("=== DCN-M2 OffsetClampSTE ===")
|
||||
offsets = torch.randn(2, 14, 16, 16, requires_grad=True) * 5.0 # outliers ~ ±15
|
||||
print(f" in: range [{offsets.min().item():.2f}, {offsets.max().item():.2f}]")
|
||||
clamped = offset_clamp_ste(offsets, kernel_size=3)
|
||||
print(f" out: range [{clamped.min().item():.2f}, {clamped.max().item():.2f}]")
|
||||
clamped.sum().backward()
|
||||
grad_mean = offsets.grad.mean().item()
|
||||
grad_std = offsets.grad.std().item()
|
||||
print(f" grad mean={grad_mean:.4f} std={grad_std:.4f} "
|
||||
f"(STE → exactly 1.0 for every element)")
|
||||
|
||||
print("\n=== KScaledFakeQuant (num_bins=3) ===")
|
||||
fq = KScaledFakeQuant(num_bins=3)
|
||||
fq.start_calibration()
|
||||
with torch.no_grad():
|
||||
for _ in range(20):
|
||||
x = torch.randn(64, 32) * 1.5
|
||||
x[0, 0] = 25.0 # outlier
|
||||
x[3, 7] = -40.0
|
||||
_ = fq(x)
|
||||
fq.finalize_calibration(percentile=99.9)
|
||||
print(f" thresholds: {fq.thresholds.tolist()}")
|
||||
print(f" scales: {[round(s, 5) for s in fq.scales.tolist()]}")
|
||||
|
||||
x_test = torch.randn(8, 32) * 3
|
||||
x_test[0, 0] = 30.0
|
||||
x_q = fq(x_test)
|
||||
err = (x_q - x_test).abs().mean().item()
|
||||
rel_err = err / x_test.abs().mean().item()
|
||||
print(f" avg abs err: {err:.4e} (relative: {rel_err:.2%})")
|
||||
|
||||
print("\n=== KScaledMamba2Block ===")
|
||||
block = KScaledMamba2Block(
|
||||
channels=128, d_state=32, headdim=32, num_bins=3,
|
||||
targets=("delta", "y"),
|
||||
)
|
||||
n_params = sum(p.numel() for p in block.parameters()) / 1e3
|
||||
print(f" params: {n_params:.1f} K")
|
||||
|
||||
# Calibration on dummy data
|
||||
block.start_calibration()
|
||||
with torch.no_grad():
|
||||
for _ in range(3):
|
||||
_ = block(torch.randn(1, 64, 128))
|
||||
block.finalize_calibration(percentile=99.9)
|
||||
|
||||
# FP vs k-scaled-INT8 comparison
|
||||
x = torch.randn(2, 64, 128)
|
||||
block.set_quant_enabled(False)
|
||||
y_fp = block(x)
|
||||
block.set_quant_enabled(True)
|
||||
y_q = block(x)
|
||||
diff = (y_fp - y_q).abs().mean().item()
|
||||
rel = diff / y_fp.abs().mean().item()
|
||||
print(f" FP output range [{y_fp.min().item():.3f}, {y_fp.max().item():.3f}]")
|
||||
print(f" INT8 output range [{y_q.min().item():.3f}, {y_q.max().item():.3f}]")
|
||||
print(f" abs diff: {diff:.4e} (relative: {rel:.2%})")
|
||||
158
src/models/sofia_v71/verify.py
Normal file
158
src/models/sofia_v71/verify.py
Normal file
@@ -0,0 +1,158 @@
|
||||
"""Verify SOFIA v7.1 model scales: param count, memory footprint, forward pass.
|
||||
|
||||
Run:
|
||||
python -m sofia_v71.verify
|
||||
python -m sofia_v71.verify --preset L
|
||||
python -m sofia_v71.verify --preset M --benchmark
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from .blocks import is_mamba_ssm_available, is_mamba2_available
|
||||
from .config import sofia_m_config, sofia_l_config, sofia_tiny_config
|
||||
from .model import SOFIAv71
|
||||
|
||||
|
||||
def count_parameters(model: torch.nn.Module) -> dict:
|
||||
"""Count parameters per named child."""
|
||||
counts = {}
|
||||
for name, param in model.named_parameters():
|
||||
top = name.split(".")[0]
|
||||
counts[top] = counts.get(top, 0) + param.numel()
|
||||
counts["_total"] = sum(counts.values())
|
||||
return counts
|
||||
|
||||
|
||||
def estimate_quantized_size(n_params: int, precision: str = "int8") -> float:
|
||||
"""Estimate on-disk / VRAM weight storage in MB."""
|
||||
bytes_per_param = {"fp32": 4, "fp16": 2, "int8": 1, "int4": 0.5}[precision]
|
||||
return n_params * bytes_per_param / (1024 ** 2)
|
||||
|
||||
|
||||
def test_forward(model: SOFIAv71, device: str = "cpu", batch_size: int = 1) -> dict:
|
||||
"""Dry forward pass, return output shapes."""
|
||||
model = model.to(device).eval()
|
||||
cfg = model.cfg
|
||||
sat = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device)
|
||||
uav = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device)
|
||||
altitude = torch.rand(batch_size, device=device) * 300.0 + 50.0
|
||||
|
||||
shapes = {}
|
||||
with torch.no_grad():
|
||||
out = model(sat=sat, uav=uav, altitude=altitude)
|
||||
for k, v in out.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
shapes[k] = tuple(v.shape)
|
||||
elif isinstance(v, list):
|
||||
shapes[k] = f"list[{len(v)} × {tuple(v[0].shape)}]"
|
||||
elif isinstance(v, dict):
|
||||
for kk, vv in v.items():
|
||||
if isinstance(vv, torch.Tensor):
|
||||
shapes[f"{k}.{kk}"] = tuple(vv.shape)
|
||||
return shapes
|
||||
|
||||
|
||||
def benchmark(model: SOFIAv71, device: str = "cpu", n_warmup: int = 3, n_runs: int = 20) -> dict:
|
||||
"""Measure forward pass latency."""
|
||||
model = model.to(device).eval()
|
||||
cfg = model.cfg
|
||||
sat = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device)
|
||||
uav = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device)
|
||||
altitude = torch.tensor([150.0], device=device)
|
||||
|
||||
# Warmup
|
||||
with torch.no_grad():
|
||||
for _ in range(n_warmup):
|
||||
_ = model(sat=sat, uav=uav, altitude=altitude)
|
||||
if device.startswith("cuda"):
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Measured runs
|
||||
times = []
|
||||
for _ in range(n_runs):
|
||||
if device.startswith("cuda"):
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
_ = model(sat=sat, uav=uav, altitude=altitude)
|
||||
if device.startswith("cuda"):
|
||||
torch.cuda.synchronize()
|
||||
times.append(time.perf_counter() - t0)
|
||||
|
||||
times_ms = [t * 1000 for t in times]
|
||||
return {
|
||||
"mean_ms": sum(times_ms) / len(times_ms),
|
||||
"min_ms": min(times_ms),
|
||||
"max_ms": max(times_ms),
|
||||
"std_ms": (sum((t - sum(times_ms) / len(times_ms)) ** 2 for t in times_ms) / len(times_ms)) ** 0.5,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Verify SOFIA v7.1 model")
|
||||
parser.add_argument("--preset", choices=["Tiny", "M", "L"], default="Tiny",
|
||||
help="Model preset (default: Tiny for fast smoke test)")
|
||||
parser.add_argument("--device", default="cpu")
|
||||
parser.add_argument("--benchmark", action="store_true")
|
||||
parser.add_argument("--batch-size", type=int, default=1)
|
||||
args = parser.parse_args()
|
||||
|
||||
cfg_fn = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}[args.preset]
|
||||
cfg = cfg_fn()
|
||||
|
||||
# Report mamba_ssm availability
|
||||
print(f"\n=== Environment ===")
|
||||
print(f" mamba_ssm v1 (Mamba-1 selective scan): "
|
||||
f"{'available' if is_mamba_ssm_available() else 'NOT available'}")
|
||||
print(f" mamba_ssm Mamba-2 (SSD dual form): "
|
||||
f"{'available' if is_mamba2_available() else 'NOT available'}")
|
||||
print(f" configured variant: {cfg.mamba_variant}, backend: {cfg.mamba_backend}")
|
||||
|
||||
print(f"\n=== SOFIA-{args.preset} configuration ===")
|
||||
print(cfg.summary())
|
||||
|
||||
print("\n=== Building model ===")
|
||||
model = SOFIAv71(cfg)
|
||||
counts = count_parameters(model)
|
||||
total = counts["_total"]
|
||||
|
||||
print(f"\nTotal params: {total / 1e6:.2f} M")
|
||||
print("\nPer-module breakdown:")
|
||||
for k, v in sorted(counts.items(), key=lambda x: -x[1]):
|
||||
if k == "_total":
|
||||
continue
|
||||
pct = 100 * v / total
|
||||
print(f" {k:30s} {v / 1e6:>8.3f} M ({pct:5.1f}%)")
|
||||
|
||||
print("\n=== Memory footprint estimates ===")
|
||||
for prec in ["fp32", "fp16", "int8", "int4"]:
|
||||
mb = estimate_quantized_size(total, prec)
|
||||
print(f" {prec.upper():6s} weights: {mb:>8.1f} MB ({mb / 1024:.2f} GB)")
|
||||
|
||||
print("\n=== Forward pass test ===")
|
||||
try:
|
||||
shapes = test_forward(model, device=args.device, batch_size=args.batch_size)
|
||||
for k, v in shapes.items():
|
||||
print(f" out[{k}]: {v}")
|
||||
except Exception as e:
|
||||
print(f" ERROR during forward: {type(e).__name__}: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
if args.benchmark:
|
||||
print("\n=== Latency benchmark ===")
|
||||
try:
|
||||
stats = benchmark(model, device=args.device)
|
||||
print(f" mean: {stats['mean_ms']:.2f} ms min: {stats['min_ms']:.2f} "
|
||||
f"max: {stats['max_ms']:.2f} std: {stats['std_ms']:.2f}")
|
||||
except Exception as e:
|
||||
print(f" ERROR during benchmark: {type(e).__name__}: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
104
src/training/csv_logger.py
Normal file
104
src/training/csv_logger.py
Normal file
@@ -0,0 +1,104 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Per-batch and per-epoch CSV logger.
|
||||
|
||||
Writes:
|
||||
{output_dir}/logs/train.csv — epoch-level train averages
|
||||
{output_dir}/logs/val.csv — epoch-level val metrics
|
||||
{output_dir}/logs/train_recall.csv — epoch-level train recall metrics
|
||||
{output_dir}/logs/train_batches.csv — per-batch train metrics (all epochs)
|
||||
{output_dir}/logs/epoch_{N}_batches.csv — per-batch for one epoch
|
||||
|
||||
Body transplanted verbatim from src/training/train_gtauav.py (pre-step-4b)
|
||||
with no logic changes — only the relocation.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.csv_logger")
|
||||
|
||||
|
||||
class CSVLogger:
|
||||
"""Log train/val metrics to CSV files using pandas."""
|
||||
|
||||
def __init__(self, output_dir: Path) -> None:
|
||||
self.log_dir = output_dir / "logs"
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
self._current_epoch: int = -1
|
||||
self._batch_columns: list[str] | None = None
|
||||
self._cumulative_batch_path = self.log_dir / "train_batches.csv"
|
||||
self._epoch_batch_path: Path | None = None
|
||||
|
||||
# Load existing CSV data on resume (so plots show full history).
|
||||
train_csv = self.log_dir / "train.csv"
|
||||
val_csv = self.log_dir / "val.csv"
|
||||
train_recall_csv = self.log_dir / "train_recall.csv"
|
||||
if train_csv.exists():
|
||||
self.train_rows = pd.read_csv(train_csv).to_dict("records")
|
||||
LOGGER.info("CSVLogger: loaded %d previous train epochs", len(self.train_rows))
|
||||
else:
|
||||
self.train_rows = []
|
||||
if val_csv.exists():
|
||||
self.val_rows = pd.read_csv(val_csv).to_dict("records")
|
||||
LOGGER.info("CSVLogger: loaded %d previous val epochs", len(self.val_rows))
|
||||
else:
|
||||
self.val_rows = []
|
||||
if train_recall_csv.exists():
|
||||
self.train_recall_rows = pd.read_csv(train_recall_csv).to_dict("records")
|
||||
else:
|
||||
self.train_recall_rows = []
|
||||
|
||||
def log_batch(self, epoch: int, batch_idx: int, global_step: int, metrics: dict) -> None:
|
||||
"""Log metrics for a single training batch. Writes to disk immediately."""
|
||||
row = {"epoch": epoch, "batch": batch_idx, "global_step": global_step, **metrics}
|
||||
|
||||
# On new epoch, start a fresh per-epoch CSV.
|
||||
if epoch != self._current_epoch:
|
||||
self._current_epoch = epoch
|
||||
self._epoch_batch_path = self.log_dir / f"epoch_{epoch:03d}_batches.csv"
|
||||
|
||||
# Determine columns on first call (consistent order).
|
||||
if self._batch_columns is None:
|
||||
self._batch_columns = list(row.keys())
|
||||
|
||||
row_df = pd.DataFrame([row], columns=self._batch_columns)
|
||||
write_header = not self._cumulative_batch_path.exists()
|
||||
|
||||
# Append to cumulative CSV.
|
||||
row_df.to_csv(
|
||||
self._cumulative_batch_path, mode="a", header=write_header, index=False,
|
||||
)
|
||||
# Append to per-epoch CSV.
|
||||
write_epoch_header = not self._epoch_batch_path.exists()
|
||||
row_df.to_csv(
|
||||
self._epoch_batch_path, mode="a", header=write_epoch_header, index=False,
|
||||
)
|
||||
|
||||
def log_train(self, epoch: int, metrics: dict, lr: float, elapsed: float) -> None:
|
||||
"""Log epoch-level train averages. Replaces existing entry for same epoch on resume."""
|
||||
row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics}
|
||||
# Remove previous entry for this epoch (resume may re-run it).
|
||||
self.train_rows = [r for r in self.train_rows if r.get("epoch") != epoch]
|
||||
self.train_rows.append(row)
|
||||
pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False)
|
||||
|
||||
def log_val(self, epoch: int, metrics: dict) -> None:
|
||||
"""Log val metrics. Replaces existing entry for same epoch on resume."""
|
||||
row = {"epoch": epoch, **metrics}
|
||||
self.val_rows = [r for r in self.val_rows if r.get("epoch") != epoch]
|
||||
self.val_rows.append(row)
|
||||
pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False)
|
||||
|
||||
def log_train_recall(self, epoch: int, metrics: dict) -> None:
|
||||
"""Log train recall metrics. Replaces existing entry for same epoch."""
|
||||
row = {"epoch": epoch, **metrics}
|
||||
self.train_recall_rows = [r for r in self.train_recall_rows if r.get("epoch") != epoch]
|
||||
self.train_recall_rows.append(row)
|
||||
pd.DataFrame(self.train_recall_rows).to_csv(
|
||||
self.log_dir / "train_recall.csv", index=False,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,273 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Training loop for caption quality test on cross-view geo-localization.
|
||||
|
||||
GeoRSCLIP dual encoder with GatedFusion on query branch.
|
||||
Single InfoNCE loss: query(drone+text) vs gallery(satellite).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import gin
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.amp import GradScaler, autocast
|
||||
from torch.optim import AdamW
|
||||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from src.datasets.visloc_with_captions import (
|
||||
GeoLocCaptionDataset,
|
||||
collate_caption_batch,
|
||||
)
|
||||
from src.eval.evaluate import evaluate_retrieval
|
||||
from src.losses.multi_infonce import InfoNCELoss
|
||||
from src.models.dual_encoder import DualEncoderCaptionTest
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.train")
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class TrainConfig:
|
||||
"""Top-level training configuration.
|
||||
|
||||
Args:
|
||||
train_query_file: Path to train_query.txt.
|
||||
val_query_file: Path to test_query.txt (used as val).
|
||||
data_root: Root of UAV-GeoLoc dataset.
|
||||
output_dir: Checkpoint and log output directory.
|
||||
epochs: Number of training epochs.
|
||||
batch_size: Mini-batch size.
|
||||
num_workers: DataLoader workers.
|
||||
learning_rate: AdamW initial LR.
|
||||
weight_decay: AdamW weight decay.
|
||||
grad_clip: Max gradient norm (0 disables).
|
||||
use_amp: Enable fp16 mixed-precision.
|
||||
eval_every: Run validation every N epochs.
|
||||
seed: Random seed.
|
||||
device: torch device.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_query_file: str = "Index/train_query.txt",
|
||||
val_query_file: str = "Index/test_query.txt",
|
||||
data_root: str = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc",
|
||||
output_dir: str = "out/caption_test",
|
||||
epochs: int = 10,
|
||||
batch_size: int = 128,
|
||||
num_workers: int = 4,
|
||||
learning_rate: float = 1e-4,
|
||||
weight_decay: float = 1e-4,
|
||||
grad_clip: float = 1.0,
|
||||
use_amp: bool = True,
|
||||
eval_every: int = 2,
|
||||
seed: int = 42,
|
||||
device: str = "cuda",
|
||||
) -> None:
|
||||
self.train_query_file = train_query_file
|
||||
self.val_query_file = val_query_file
|
||||
self.data_root = data_root
|
||||
self.output_dir = Path(output_dir)
|
||||
self.epochs = epochs
|
||||
self.batch_size = batch_size
|
||||
self.num_workers = num_workers
|
||||
self.learning_rate = learning_rate
|
||||
self.weight_decay = weight_decay
|
||||
self.grad_clip = grad_clip
|
||||
self.use_amp = use_amp
|
||||
self.eval_every = eval_every
|
||||
self.seed = seed
|
||||
self.device = device
|
||||
|
||||
|
||||
def _set_seed(seed: int) -> None:
|
||||
import random as _random
|
||||
import numpy as _np
|
||||
_random.seed(seed)
|
||||
_np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def _atomic_save(obj: dict, path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = path.with_suffix(path.suffix + ".tmp")
|
||||
torch.save(obj, tmp_path)
|
||||
tmp_path.replace(path)
|
||||
|
||||
|
||||
def train(config_path: str) -> None:
|
||||
"""Run full training loop from gin config."""
|
||||
gin.parse_config_file(config_path)
|
||||
cfg = TrainConfig()
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||||
)
|
||||
_set_seed(cfg.seed)
|
||||
cfg.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Model + loss.
|
||||
model = DualEncoderCaptionTest().to(cfg.device)
|
||||
loss_fn = InfoNCELoss().to(cfg.device)
|
||||
|
||||
preprocess = model.preprocess
|
||||
|
||||
train_ds = GeoLocCaptionDataset(
|
||||
query_file=cfg.train_query_file,
|
||||
data_root=cfg.data_root,
|
||||
image_transform=preprocess,
|
||||
)
|
||||
val_ds = GeoLocCaptionDataset(
|
||||
query_file=cfg.val_query_file,
|
||||
data_root=cfg.data_root,
|
||||
image_transform=preprocess,
|
||||
)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
batch_size=cfg.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=cfg.num_workers,
|
||||
collate_fn=collate_caption_batch,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_ds,
|
||||
batch_size=cfg.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=cfg.num_workers,
|
||||
collate_fn=collate_caption_batch,
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
optimizer = AdamW(
|
||||
model.trainable_parameters(),
|
||||
lr=cfg.learning_rate,
|
||||
weight_decay=cfg.weight_decay,
|
||||
)
|
||||
scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs)
|
||||
scaler = GradScaler(enabled=cfg.use_amp)
|
||||
|
||||
n_trainable = sum(p.numel() for p in model.trainable_parameters())
|
||||
n_total = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info(
|
||||
"trainable=%d (%.2f%%) total=%d train=%d val=%d",
|
||||
n_trainable, 100.0 * n_trainable / n_total,
|
||||
n_total, len(train_ds), len(val_ds),
|
||||
)
|
||||
|
||||
history: list[dict] = []
|
||||
|
||||
for epoch in range(cfg.epochs):
|
||||
model.train()
|
||||
epoch_start = time.time()
|
||||
agg: dict[str, float] = {}
|
||||
n_batches = 0
|
||||
|
||||
for batch in train_loader:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
|
||||
sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
|
||||
caption_drone = batch["caption_drone"]
|
||||
|
||||
with autocast(device_type="cuda", enabled=cfg.use_amp):
|
||||
embeddings = model(
|
||||
drone_img=drone_img,
|
||||
sat_img=sat_img,
|
||||
caption_drone=caption_drone,
|
||||
)
|
||||
loss_dict = loss_fn(
|
||||
embeddings=embeddings,
|
||||
epoch=epoch,
|
||||
total_epochs=cfg.epochs,
|
||||
)
|
||||
|
||||
total_loss = loss_dict["total"]
|
||||
scaler.scale(total_loss).backward()
|
||||
|
||||
if cfg.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
nn.utils.clip_grad_norm_(
|
||||
model.trainable_parameters(),
|
||||
max_norm=cfg.grad_clip,
|
||||
)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
|
||||
for key, val in loss_dict.items():
|
||||
agg[key] = agg.get(key, 0.0) + float(val.item())
|
||||
n_batches += 1
|
||||
|
||||
scheduler.step()
|
||||
elapsed = time.time() - epoch_start
|
||||
|
||||
means = {k: v / max(n_batches, 1) for k, v in agg.items()}
|
||||
LOGGER.info(
|
||||
"epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate=%.4f",
|
||||
epoch, elapsed,
|
||||
optimizer.param_groups[0]["lr"],
|
||||
means.get("total", 0.0),
|
||||
means.get("temperature", 0.0),
|
||||
means.get("gate", 1.0),
|
||||
)
|
||||
|
||||
epoch_record: dict = {
|
||||
"epoch": epoch,
|
||||
"elapsed_seconds": elapsed,
|
||||
"train": means,
|
||||
}
|
||||
|
||||
# Validation.
|
||||
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
|
||||
model.eval()
|
||||
val_metrics = evaluate_retrieval(
|
||||
model=model,
|
||||
loader=val_loader,
|
||||
device=cfg.device,
|
||||
)
|
||||
epoch_record["val"] = val_metrics
|
||||
LOGGER.info(
|
||||
"val epoch=%d R@1_q2g=%.4f R@5_q2g=%.4f R@10_q2g=%.4f",
|
||||
epoch,
|
||||
val_metrics.get("r@1_query_to_gallery", 0.0),
|
||||
val_metrics.get("r@5_query_to_gallery", 0.0),
|
||||
val_metrics.get("r@10_query_to_gallery", 0.0),
|
||||
)
|
||||
|
||||
history.append(epoch_record)
|
||||
|
||||
_atomic_save(
|
||||
obj={
|
||||
"epoch": epoch,
|
||||
"model_state": model.state_dict(),
|
||||
"optimizer_state": optimizer.state_dict(),
|
||||
"config_path": config_path,
|
||||
},
|
||||
path=cfg.output_dir / f"ckpt_epoch{epoch:03d}.pt",
|
||||
)
|
||||
|
||||
history_path = cfg.output_dir / "history.json"
|
||||
with history_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(history, f, indent=2)
|
||||
|
||||
LOGGER.info("training complete, history saved to %s", history_path)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Caption quality test training.")
|
||||
parser.add_argument("--config", type=str, required=True, help="Gin config file.")
|
||||
args = parser.parse_args()
|
||||
train(config_path=args.config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
1060
src/training/trainer_new.py
Normal file
1060
src/training/trainer_new.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)
|
||||
|
||||
|
||||
|
||||
157
tests/conftest.py
Normal file
157
tests/conftest.py
Normal file
@@ -0,0 +1,157 @@
|
||||
"""Shared pytest fixtures + reporter hooks for caption-test test suite.
|
||||
|
||||
Provides:
|
||||
- proj_dir, path2cfg: real repository paths for integration-style tests
|
||||
that load actual gin presets from in/config_files/.
|
||||
- clear_gin: autouse fixture that wipes gin global state before every test
|
||||
(gin keeps bindings in module-level singleton; tests must not leak).
|
||||
|
||||
Reporter hooks print "✅/❌ <docstring>" next to each test's PASSED/FAILED
|
||||
line in -v mode.
|
||||
|
||||
Preset name lists live in a separate module (`tests/_presets.py`) so test
|
||||
files can import them with a plain `import _presets` — no relative imports,
|
||||
no need for tests/ to be a package.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import gin
|
||||
import pytest
|
||||
|
||||
|
||||
DINOV3_PRESETS = (
|
||||
"gtauav_balanced",
|
||||
"gtauav_balanced_asym",
|
||||
"gtauav_baseline",
|
||||
"gtauav_baseline_asym",
|
||||
"gtauav_image_heavy",
|
||||
"gtauav_text_heavy",
|
||||
)
|
||||
|
||||
STRIPNET_PRESETS = (
|
||||
"gtauav_balanced_stripnet",
|
||||
"gtauav_balanced_stripnet_unfrozen",
|
||||
"gtauav_baseline_stripnet",
|
||||
"gtauav_baseline_stripnet_unfrozen",
|
||||
)
|
||||
|
||||
ALL_TRAINING_PRESETS = DINOV3_PRESETS + STRIPNET_PRESETS
|
||||
|
||||
# --- gin hygiene -----------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_gin():
|
||||
"""Wipe gin's global binding state before AND after each test.
|
||||
|
||||
Gin keeps bindings in module-level singletons; without this fixture, a
|
||||
test that loads a config (or even just calls gin.parse_config_file in a
|
||||
helper) leaks bindings into the next test, leading to flaky failures
|
||||
like 'Unknown configurable' or wrong field values.
|
||||
"""
|
||||
gin.clear_config()
|
||||
yield
|
||||
gin.clear_config()
|
||||
|
||||
|
||||
# --- real-repo paths -------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proj_dir() -> Path:
|
||||
"""Path to repository root (the directory that contains src/, tests/, in/)."""
|
||||
return Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def path2cfg(proj_dir: Path) -> str:
|
||||
"""Trailing-slashed path to in/config_files/, matching `src/main.py`.
|
||||
|
||||
Per REQUIREMENTS_GIN_STYLE.md §5, src/main.py builds this path as
|
||||
`f"{proj_dir}in/config_files/"`. Tests that exercise the real repo
|
||||
layout should use this fixture verbatim instead of constructing it
|
||||
independently.
|
||||
"""
|
||||
return f"{proj_dir}/in/config_files/"
|
||||
|
||||
|
||||
# --- reporter hooks --------------------------------------------------------
|
||||
|
||||
|
||||
def _docstring_summary(item: pytest.Item) -> str | None:
|
||||
"""Return the first non-empty line of a test's docstring, or None."""
|
||||
func = getattr(item, "function", None) or getattr(item, "obj", None)
|
||||
if func is None or not getattr(func, "__doc__", None):
|
||||
return None
|
||||
for line in func.__doc__.strip().splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
return stripped
|
||||
return None
|
||||
|
||||
|
||||
# Cache nodeid → docstring summary, populated at collection time so the
|
||||
# logreport hook can look them up without re-introspecting the test function.
|
||||
_LAST_SEEN_SUMMARY: dict[str, str] = {}
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(
|
||||
config: pytest.Config,
|
||||
items: list[pytest.Item],
|
||||
) -> None:
|
||||
"""Cache each item's docstring summary for later use by the status hook."""
|
||||
for item in items:
|
||||
summary = _docstring_summary(item)
|
||||
if summary:
|
||||
_LAST_SEEN_SUMMARY[item.nodeid] = summary
|
||||
|
||||
|
||||
def pytest_runtest_logreport(report: pytest.TestReport) -> None:
|
||||
"""Print test results with parametrized tests grouped under one header.
|
||||
|
||||
Non-parametrized test:
|
||||
✅ <docstring summary>
|
||||
|
||||
Parametrized test (first occurrence of the group):
|
||||
<docstring summary>
|
||||
✅ <param>
|
||||
|
||||
Parametrized test (subsequent occurrences):
|
||||
✅ <param>
|
||||
|
||||
Pytest emits 3 reports per test (setup → call → teardown). We hook the
|
||||
`call` phase — the one where pass/fail is actually decided.
|
||||
"""
|
||||
if report.when != "call":
|
||||
return
|
||||
|
||||
summary = _LAST_SEEN_SUMMARY.get(report.nodeid, "(no docstring)")
|
||||
if report.passed:
|
||||
icon = "✅"
|
||||
elif report.failed:
|
||||
icon = "❌"
|
||||
else:
|
||||
icon = "⏭️"
|
||||
|
||||
# Detect parametrization: pytest encodes params as `nodeid[param1-param2-...]`.
|
||||
if "[" in report.nodeid and report.nodeid.endswith("]"):
|
||||
base_id, _, param_part = report.nodeid.partition("[")
|
||||
param_label = param_part[:-1] or "empty" # strip trailing ']'
|
||||
|
||||
# Print docstring header only on first encounter of this group.
|
||||
# Leading "\n" separates the header from pytest's progress dot.
|
||||
if base_id not in _PRINTED_GROUP_HEADERS:
|
||||
print(f"\n{summary}")
|
||||
_PRINTED_GROUP_HEADERS.add(base_id)
|
||||
print(f" {icon} {param_label}")
|
||||
else:
|
||||
print(f" {icon} {summary}")
|
||||
|
||||
|
||||
# Tracks which parametrized test groups have already had their docstring
|
||||
# header printed. Reset implicitly at the start of each pytest run because
|
||||
# Python re-imports conftest.py.
|
||||
_PRINTED_GROUP_HEADERS: set[str] = set()
|
||||
187
tests/test_trainer.py
Normal file
187
tests/test_trainer.py
Normal file
@@ -0,0 +1,187 @@
|
||||
"""Tests for src.training.trainer_new.Trainer.
|
||||
|
||||
Scope: __init__ behaviour, backbone validation, ModelsConfig type union.
|
||||
Out of scope: actual training (requires GPU + datasets + model checkpoints).
|
||||
|
||||
The Trainer class is designed to defer all heavy lifting (CUDA, model
|
||||
construction, dataset loading) to .train(); __init__ just stores the 6 cfg
|
||||
objects and zeros out runtime state. This makes it cheap to test.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import get_args
|
||||
|
||||
import pytest
|
||||
|
||||
from src.conf.config_loader import load_all_configs
|
||||
from src.conf.models_dinov3_conf import DINOv3ModelsConfig
|
||||
from src.conf.models_stripnet_conf import StripNetModelsConfig
|
||||
from src.training.trainer_new import (
|
||||
ModelsConfig,
|
||||
Trainer,
|
||||
_SUPPORTED_BACKBONES,
|
||||
)
|
||||
|
||||
from conftest import DINOV3_PRESETS, STRIPNET_PRESETS
|
||||
|
||||
|
||||
# -- module-level constants ------------------------------------------------
|
||||
|
||||
|
||||
def test_supported_backbones_is_frozenset() -> None:
|
||||
"""_SUPPORTED_BACKBONES must be a frozenset (immutable, hashable)."""
|
||||
assert isinstance(_SUPPORTED_BACKBONES, frozenset)
|
||||
|
||||
|
||||
def test_supported_backbones_contents() -> None:
|
||||
"""Exactly dinov3 and stripnet are supported in the current refactor.
|
||||
|
||||
Sofia (v1/v71) is intentionally absent — see Trainer._validate_backbone
|
||||
for the rationale and the steps to add it later.
|
||||
"""
|
||||
assert _SUPPORTED_BACKBONES == frozenset({"dinov3", "stripnet"})
|
||||
|
||||
|
||||
def test_models_config_union_contents() -> None:
|
||||
"""ModelsConfig union mirrors _SUPPORTED_BACKBONES (dinov3 | stripnet)."""
|
||||
union_members = set(get_args(ModelsConfig))
|
||||
assert union_members == {DINOv3ModelsConfig, StripNetModelsConfig}
|
||||
|
||||
|
||||
# -- _validate_backbone --------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backbone", ["dinov3", "stripnet"])
|
||||
def test_validate_backbone_accepts_supported(
|
||||
path2cfg: str, backbone: str,
|
||||
) -> None:
|
||||
"""Supported backbones pass _validate_backbone silently.
|
||||
|
||||
We use a real preset to build a valid Trainer — this also exercises
|
||||
the load_all_configs → Trainer(...) integration.
|
||||
"""
|
||||
preset = "gtauav_balanced" if backbone == "dinov3" else "gtauav_balanced_stripnet"
|
||||
cfgs = load_all_configs(path2cfg, preset)
|
||||
trainer = Trainer(
|
||||
pipeline_cfg=cfgs["pipeline"],
|
||||
hardware_cfg=cfgs["hardware"],
|
||||
training_cfg=cfgs["training"],
|
||||
tracking_cfg=cfgs["tracking"],
|
||||
models_common_cfg=cfgs["models_common"],
|
||||
models_cfg=cfgs["models"],
|
||||
)
|
||||
# Must not raise.
|
||||
trainer._validate_backbone()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_backbone", ["sofia_v1", "sofia_v71", "mistral_42b", ""])
|
||||
def test_validate_backbone_rejects_unsupported(
|
||||
path2cfg: str, bad_backbone: str,
|
||||
) -> None:
|
||||
"""Unsupported backbones (incl. sofia) raise NotImplementedError, not ImportError.
|
||||
|
||||
The user must get a clear, actionable message — not a stack trace from
|
||||
a missing module.
|
||||
"""
|
||||
cfgs = load_all_configs(path2cfg, "gtauav_balanced")
|
||||
trainer = Trainer(
|
||||
pipeline_cfg=cfgs["pipeline"],
|
||||
hardware_cfg=cfgs["hardware"],
|
||||
training_cfg=cfgs["training"],
|
||||
tracking_cfg=cfgs["tracking"],
|
||||
models_common_cfg=cfgs["models_common"],
|
||||
models_cfg=cfgs["models"],
|
||||
)
|
||||
# Tamper with backbone — simulate what would happen if config_loader
|
||||
# were extended to accept sofia presets.
|
||||
trainer.models_common_cfg.backbone = bad_backbone
|
||||
|
||||
with pytest.raises(NotImplementedError) as excinfo:
|
||||
trainer._validate_backbone()
|
||||
|
||||
# Error message must mention the offending backbone name and what's supported.
|
||||
msg = str(excinfo.value)
|
||||
assert bad_backbone in msg or repr(bad_backbone) in msg
|
||||
assert "dinov3" in msg
|
||||
assert "stripnet" in msg
|
||||
|
||||
|
||||
# -- Trainer.__init__ smoke tests ------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preset_name", DINOV3_PRESETS + STRIPNET_PRESETS)
|
||||
def test_trainer_init_with_real_preset(path2cfg: str, preset_name: str) -> None:
|
||||
"""Trainer(...) instantiates from every real preset's loaded cfgs.
|
||||
|
||||
Heavy work (CUDA, model build, dataset open) is deferred to .train();
|
||||
__init__ only stores cfgs and zeros runtime state, so this is cheap and
|
||||
GPU-free.
|
||||
"""
|
||||
cfgs = load_all_configs(path2cfg, preset_name)
|
||||
|
||||
trainer = Trainer(
|
||||
pipeline_cfg=cfgs["pipeline"],
|
||||
hardware_cfg=cfgs["hardware"],
|
||||
training_cfg=cfgs["training"],
|
||||
tracking_cfg=cfgs["tracking"],
|
||||
models_common_cfg=cfgs["models_common"],
|
||||
models_cfg=cfgs["models"],
|
||||
)
|
||||
|
||||
# Cfgs are stored as-is.
|
||||
assert trainer.pipeline_cfg is cfgs["pipeline"]
|
||||
assert trainer.hardware_cfg is cfgs["hardware"]
|
||||
assert trainer.training_cfg is cfgs["training"]
|
||||
assert trainer.tracking_cfg is cfgs["tracking"]
|
||||
assert trainer.models_common_cfg is cfgs["models_common"]
|
||||
assert trainer.models_cfg is cfgs["models"]
|
||||
|
||||
|
||||
def test_trainer_init_zeros_runtime_state(path2cfg: str) -> None:
|
||||
"""All runtime fields are None / 0 / [] before .train() is called."""
|
||||
cfgs = load_all_configs(path2cfg, "gtauav_balanced")
|
||||
trainer = Trainer(
|
||||
pipeline_cfg=cfgs["pipeline"],
|
||||
hardware_cfg=cfgs["hardware"],
|
||||
training_cfg=cfgs["training"],
|
||||
tracking_cfg=cfgs["tracking"],
|
||||
models_common_cfg=cfgs["models_common"],
|
||||
models_cfg=cfgs["models"],
|
||||
)
|
||||
|
||||
# None-typed runtime fields.
|
||||
for attr in (
|
||||
"output_dir", "full_config", "tracker", "csv_logger", "model",
|
||||
"loss_fn", "neg_bank", "optimizer", "scheduler", "scaler",
|
||||
"train_ds", "test_ds", "train_eval_ds",
|
||||
"train_loader", "test_loader", "train_eval_loader",
|
||||
"batch_sampler", "emb_cache", "profiler", "resume_ckpt",
|
||||
):
|
||||
assert getattr(trainer, attr) is None, (
|
||||
f"trainer.{attr} should be None before .train(), "
|
||||
f"got {type(getattr(trainer, attr)).__name__}"
|
||||
)
|
||||
|
||||
# Counter / loop state initialized to identity values.
|
||||
assert trainer.start_epoch == 0
|
||||
assert trainer.global_step == 0
|
||||
assert trainer.best_r1 == 0.0
|
||||
assert trainer.history == []
|
||||
assert trainer.steps_per_epoch == 0
|
||||
|
||||
|
||||
# -- Trainer.train end-to-end signature ------------------------------------
|
||||
|
||||
|
||||
def test_trainer_train_method_exists_and_takes_no_args() -> None:
|
||||
"""Trainer.train() takes only `self` — main.py calls trainer.train()."""
|
||||
import inspect
|
||||
|
||||
sig = inspect.signature(Trainer.train)
|
||||
params = [p for p in sig.parameters.values() if p.name != "self"]
|
||||
assert params == [], (
|
||||
f"Trainer.train() must take only self; got extra params: {params}"
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user