Compare commits
2 Commits
belka_exp_
...
belka_refa
| Author | SHA1 | Date | |
|---|---|---|---|
|
0d8c82acc3 | ||
|
|
a27b5a7357 |
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-репозиторий авторов
|
||||
**Путь к данным:** `/home/servml/Документы/datasets/GTA-UAV-LR/`
|
||||
**Путь к аугментациям:** `/home/servml/Документы/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)
|
||||
|
||||
**Путь:** `/home/servml/Документы/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/
|
||||
```
|
||||
@@ -34,10 +34,10 @@ TrainConfigGTAUAV.label_smoothing = 0.1
|
||||
TrainConfigGTAUAV.learnable_temperature = True
|
||||
TrainConfigGTAUAV.weight_q2g = 0.6
|
||||
TrainConfigGTAUAV.weight_g2q = 0.4
|
||||
TrainConfigGTAUAV.neg_bank_size = 4096
|
||||
TrainConfigGTAUAV.neg_bank_size = 0 # 4096
|
||||
|
||||
# ---- Sampling ----
|
||||
TrainConfigGTAUAV.sampler_type = "dss" # "dss" or "mutex"
|
||||
TrainConfigGTAUAV.sampler_type = "mutex" # "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
|
||||
@@ -48,7 +48,7 @@ TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"
|
||||
# ---- Tracking ----
|
||||
TrainConfigGTAUAV.use_wandb = False
|
||||
TrainConfigGTAUAV.use_tb = True
|
||||
TrainConfigGTAUAV.use_gradcam = True
|
||||
TrainConfigGTAUAV.use_gradcam = False
|
||||
TrainConfigGTAUAV.gradcam_every = 5
|
||||
TrainConfigGTAUAV.use_profiler = False
|
||||
TrainConfigGTAUAV.log_grad_norms = True
|
||||
@@ -61,4 +61,4 @@ 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)
|
||||
InfoNCELoss.hard_mining_k = 0 # 512 # 0 = use whole queue (disable mining)
|
||||
|
||||
39
conf/gtauav_balanced_sofia.gin
Normal file
39
conf/gtauav_balanced_sofia.gin
Normal file
@@ -0,0 +1,39 @@
|
||||
# GTA-UAV Balanced (SOFIA-Tiny backbone): SOFIA v7.1 student trained from scratch
|
||||
# с двухуровневой text fusion:
|
||||
# 1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует feature map перед GGeM/CHP).
|
||||
# 2. Late-level: GatedFusion на дескрипторах (как в DINOv3/StripNet вариантах).
|
||||
#
|
||||
# Trainable (~5-7M):
|
||||
# - SOFIA backbone (Tiny, ~5M, from scratch — нет pretrained)
|
||||
# - SOFIA heads (SatHead GGeM+BN+Linear, UAVHead AltitudeFiLM+CHP+BN+Linear, +Text-FiLM)
|
||||
# - DGTRS-CLIP LoRA (rank=4, ~147K)
|
||||
# - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared)
|
||||
# - Gates α_q, α_g + learnable τ
|
||||
#
|
||||
# Altitude (drone_height метры) подаётся в UAVHead.AltitudeFiLM из dataset meta CSV.
|
||||
# Для sat — altitude=None → FiLM passthrough (γ=1, β=0).
|
||||
#
|
||||
# Note: SOFIA from scratch — нужно больше эпох или warmup, чем frozen DINOv3/StripNet.
|
||||
# Mamba-2 backend (mamba_ssm) даёт 2-8x speedup vs torch fallback.
|
||||
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
# ---- Backbone ----
|
||||
TrainConfigGTAUAV.backbone = "sofia"
|
||||
TrainConfigGTAUAV.sofia_preset = "Tiny"
|
||||
TrainConfigGTAUAV.sofia_d_descriptor = 1024
|
||||
TrainConfigGTAUAV.sofia_use_text_film_uav = True
|
||||
TrainConfigGTAUAV.sofia_use_text_film_sat = True
|
||||
TrainConfigGTAUAV.sofia_lora_rank = 4
|
||||
# Mamba-1 used for Tiny (Mamba-2 torch fallback has a pre-existing reshape bug
|
||||
# with channels not divisible by default headdim; switch to "mamba2" for M/L
|
||||
# presets where channels % 64 == 0 OR install mamba_ssm CUDA kernels).
|
||||
TrainConfigGTAUAV.sofia_mamba_variant = "mamba1"
|
||||
TrainConfigGTAUAV.sofia_mamba_backend = "auto" # mamba_ssm if installed else torch fallback
|
||||
|
||||
# ---- Training overrides ----
|
||||
TrainConfigGTAUAV.gradient_checkpointing = False # SOFIA from-scratch — keep activations live
|
||||
TrainConfigGTAUAV.shared_encoder = True # ignored by SOFIA but kept for logging compat
|
||||
|
||||
# ---- Output ----
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia"
|
||||
42
conf/gtauav_balanced_sofia_v1.gin
Normal file
42
conf/gtauav_balanced_sofia_v1.gin
Normal file
@@ -0,0 +1,42 @@
|
||||
# GTA-UAV Balanced (SOFIA v1 backbone): StripNet+DCNv4 hierarchical CNN
|
||||
# (~3-30M params depending on variant) trained from scratch с двухуровневой
|
||||
# text fusion:
|
||||
# 1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует [B,C,8,8] перед GGeM).
|
||||
# 2. Late-level: GatedFusion на дескрипторах.
|
||||
#
|
||||
# UAV head: AltitudeFiLM(drone_height) + [TextFiLM] + GGeM + BN + Linear.
|
||||
# SAT head: [TextFiLM] + GGeM + BN + Linear.
|
||||
# Один backbone shared между sat/uav.
|
||||
#
|
||||
# Variant -> размер модели:
|
||||
# tiny_tiny: dims [16, 32, 80, 128] (~0.4M)
|
||||
# tiny : dims [32, 64, 128, 256] (~1M)
|
||||
# small : dims [64, 128, 320, 512] (~3M, default)
|
||||
# small_v2 : dims [64, 128, 256, 384] (~2M)
|
||||
#
|
||||
# Trainable (с small variant + text fusion):
|
||||
# - SOFIA v1 backbone (~3M) + heads (~0.6M)
|
||||
# - DGTRS-CLIP LoRA (rank=4, ~147K)
|
||||
# - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared)
|
||||
# - Gates α_q, α_g + learnable τ
|
||||
# Total trainable ~7M.
|
||||
#
|
||||
# Note: DCNv4 требует CUDA — обучение только на GPU. Не работает на CPU.
|
||||
|
||||
include 'conf/gtauav_balanced.gin'
|
||||
|
||||
# ---- Backbone ----
|
||||
TrainConfigGTAUAV.backbone = "sofia_v1"
|
||||
TrainConfigGTAUAV.sofia_v1_variant = "tiny"
|
||||
TrainConfigGTAUAV.sofia_v1_d_descriptor = 1024
|
||||
TrainConfigGTAUAV.sofia_v1_use_text_film_uav = True
|
||||
TrainConfigGTAUAV.sofia_v1_use_text_film_sat = True
|
||||
TrainConfigGTAUAV.sofia_v1_use_film_altitude = True
|
||||
TrainConfigGTAUAV.sofia_v1_lora_rank = 4
|
||||
|
||||
# ---- Training overrides ----
|
||||
TrainConfigGTAUAV.gradient_checkpointing = False # SOFIA v1 from-scratch
|
||||
TrainConfigGTAUAV.shared_encoder = True
|
||||
|
||||
# ---- Output ----
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia_v1"
|
||||
13
conf/gtauav_baseline_sofia.gin
Normal file
13
conf/gtauav_baseline_sofia.gin
Normal file
@@ -0,0 +1,13 @@
|
||||
# GTA-UAV Baseline (SOFIA-Tiny backbone): no text fusion. Reference R@1 для
|
||||
# computing Δ R@1 vs gtauav_balanced_sofia.gin.
|
||||
#
|
||||
# В baseline_mode=True:
|
||||
# - Text-FiLM отключается (SOFIA heads работают только с altitude).
|
||||
# - DGTRS-CLIP не загружается, TextFusionMLP не строится.
|
||||
# - GatedFusion gates = 1.0 (text игнорируется).
|
||||
|
||||
include 'conf/gtauav_balanced_sofia.gin'
|
||||
|
||||
TrainConfigGTAUAV.baseline_mode = True
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia"
|
||||
TrainConfigGTAUAV.use_gradcam = False
|
||||
8
conf/gtauav_baseline_sofia_v1.gin
Normal file
8
conf/gtauav_baseline_sofia_v1.gin
Normal file
@@ -0,0 +1,8 @@
|
||||
# GTA-UAV Baseline (SOFIA v1 backbone): no text fusion. Reference R@1 для
|
||||
# computing Δ R@1 vs gtauav_balanced_sofia_v1.gin.
|
||||
|
||||
include 'conf/gtauav_balanced_sofia_v1.gin'
|
||||
|
||||
TrainConfigGTAUAV.baseline_mode = True
|
||||
TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia_v1"
|
||||
TrainConfigGTAUAV.use_gradcam = False
|
||||
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')
|
||||
@@ -104,6 +104,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 +121,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 +136,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 +149,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 +245,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 +313,7 @@ class GTAUAVDataset(Dataset):
|
||||
"pair_id": entry["drone_name"],
|
||||
"sat_name": sat_name,
|
||||
"positive_weight": pos_weight,
|
||||
"altitude": float(entry["altitude"]),
|
||||
}
|
||||
|
||||
|
||||
@@ -286,6 +333,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 +419,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 +441,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),
|
||||
}
|
||||
|
||||
@@ -462,6 +462,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 +493,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()
|
||||
@@ -54,6 +54,14 @@ from src.models.asymmetric_encoder import (
|
||||
get_drone_train_transform,
|
||||
get_satellite_train_transform,
|
||||
)
|
||||
from src.models.sofia_fusion_encoder import SOFIAFusionEncoder
|
||||
from src.models.sofia_v1 import SOFIAv1Config
|
||||
from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder
|
||||
from src.models.sofia_v71 import (
|
||||
sofia_l_config,
|
||||
sofia_m_config,
|
||||
sofia_tiny_config,
|
||||
)
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.train_gtauav")
|
||||
|
||||
@@ -91,11 +99,30 @@ class TrainConfigGTAUAV:
|
||||
mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks
|
||||
gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch)
|
||||
# StripNet backbone option (replaces DINOv3 when backbone="stripnet").
|
||||
backbone: str = "dinov3" # "dinov3" or "stripnet"
|
||||
backbone: str = "dinov3" # "dinov3", "stripnet", or "sofia"
|
||||
stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth"
|
||||
stripnet_mona_last_n_stages: int = 2 # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA)
|
||||
stripnet_freeze: bool = True # If False, StripNet backbone is fully trainable (full fine-tune)
|
||||
stripnet_backbone_lr_factor: float = 0.1 # Backbone LR = learning_rate * factor (only when unfrozen)
|
||||
# SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint.
|
||||
sofia_preset: str = "Tiny" # "Tiny" | "M" | "L"
|
||||
sofia_d_descriptor: int = 1024 # retrieval space (1024 = match TextFusionMLP out_dim)
|
||||
sofia_use_text_film_uav: bool = True # mid-level text-FiLM in UAV head
|
||||
sofia_use_text_film_sat: bool = True # mid-level text-FiLM in SAT head
|
||||
sofia_lora_rank: int = 4
|
||||
sofia_mamba_variant: str = "mamba2" # "mamba1" | "mamba2" | "efficient_vmamba"
|
||||
sofia_mamba_backend: str = "auto" # "auto" | "torch" | "mamba_ssm"
|
||||
# EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in).
|
||||
sofia_use_evss_bridge: bool = False
|
||||
sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"])
|
||||
# SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch.
|
||||
sofia_v1_variant: str = "small" # "tiny_tiny" | "tiny" | "small" | "small_v2"
|
||||
sofia_v1_dcn_variant: str = "v2" # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky)
|
||||
sofia_v1_d_descriptor: int = 1024
|
||||
sofia_v1_use_text_film_uav: bool = True
|
||||
sofia_v1_use_text_film_sat: bool = True
|
||||
sofia_v1_use_film_altitude: bool = True
|
||||
sofia_v1_lora_rank: int = 4
|
||||
|
||||
# Training.
|
||||
resume_from: str | None = None # path to checkpoint for resuming
|
||||
@@ -167,7 +194,7 @@ def _atomic_save(obj: dict, path: Path) -> None:
|
||||
|
||||
|
||||
def _build_param_groups(
|
||||
model: AsymmetricEncoder,
|
||||
model: nn.Module,
|
||||
lr: float,
|
||||
text_lr_factor: float,
|
||||
stripnet_backbone_lr_factor: float = 0.1,
|
||||
@@ -177,7 +204,8 @@ def _build_param_groups(
|
||||
Groups:
|
||||
- text_encoder.* → lr * text_lr_factor (default 1e-5)
|
||||
- image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5)
|
||||
- everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv) → lr
|
||||
- everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv,
|
||||
SOFIA backbone+heads when backbone="sofia") → lr
|
||||
"""
|
||||
text_params = []
|
||||
backbone_params = []
|
||||
@@ -249,9 +277,13 @@ def _embed_drone_queries(
|
||||
embs: list[torch.Tensor] = []
|
||||
for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False):
|
||||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||||
altitude = batch.get("altitude")
|
||||
if altitude is not None:
|
||||
altitude = altitude.to(device, non_blocking=True)
|
||||
q = model.encode_query(
|
||||
drone_img,
|
||||
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
|
||||
altitude=altitude,
|
||||
)
|
||||
embs.append(q.cpu())
|
||||
|
||||
@@ -336,9 +368,13 @@ def _evaluate(
|
||||
if max_batches is not None and i >= max_batches:
|
||||
break
|
||||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||||
altitude = batch.get("altitude")
|
||||
if altitude is not None:
|
||||
altitude = altitude.to(device, non_blocking=True)
|
||||
q = model.encode_query(
|
||||
drone_img,
|
||||
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
|
||||
altitude=altitude,
|
||||
)
|
||||
query_embs.append(q.cpu())
|
||||
query_valid_names.extend(batch["valid_sat_names"])
|
||||
@@ -575,40 +611,97 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
|
||||
if cfg.resume_from is not None:
|
||||
LOGGER.info("Resuming from %s", cfg.resume_from)
|
||||
model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
|
||||
cfg.resume_from,
|
||||
dino_web_path=cfg.dino_web_path,
|
||||
dino_sat_path=cfg.dino_sat_path,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
device=cfg.device,
|
||||
)
|
||||
if cfg.backbone == "sofia":
|
||||
model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint(
|
||||
cfg.resume_from,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
device=cfg.device,
|
||||
)
|
||||
elif cfg.backbone == "sofia_v1":
|
||||
model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint(
|
||||
cfg.resume_from,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
device=cfg.device,
|
||||
)
|
||||
else:
|
||||
model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
|
||||
cfg.resume_from,
|
||||
dino_web_path=cfg.dino_web_path,
|
||||
dino_sat_path=cfg.dino_sat_path,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
device=cfg.device,
|
||||
)
|
||||
start_epoch = resume_ckpt.get("epoch", -1) + 1
|
||||
else:
|
||||
mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
|
||||
if cfg.backbone == "stripnet":
|
||||
if cfg.backbone == "sofia":
|
||||
enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})"
|
||||
elif cfg.backbone == "sofia_v1":
|
||||
enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})"
|
||||
elif cfg.backbone == "stripnet":
|
||||
enc_str = "StripNet-small (shared, 512→1024 proj)"
|
||||
else:
|
||||
enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
|
||||
LOGGER.info("Building model — %s, %s", mode_str, enc_str)
|
||||
model = AsymmetricEncoder(
|
||||
dino_web_path=cfg.dino_web_path,
|
||||
dino_sat_path=cfg.dino_sat_path,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
init_gate=cfg.init_gate,
|
||||
baseline_mode=cfg.baseline_mode,
|
||||
shared_encoder=cfg.shared_encoder,
|
||||
mona_bottleneck=cfg.mona_bottleneck,
|
||||
mona_last_n_blocks=cfg.mona_last_n_blocks,
|
||||
device=cfg.device,
|
||||
backbone=cfg.backbone,
|
||||
stripnet_path=cfg.stripnet_path,
|
||||
stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages,
|
||||
stripnet_freeze=cfg.stripnet_freeze,
|
||||
).to(cfg.device)
|
||||
if cfg.backbone == "sofia":
|
||||
preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}
|
||||
if cfg.sofia_preset not in preset_map:
|
||||
raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}")
|
||||
sofia_cfg = preset_map[cfg.sofia_preset]()
|
||||
sofia_cfg.d_descriptor = cfg.sofia_d_descriptor
|
||||
sofia_cfg.text_film_dim = cfg.sofia_d_descriptor
|
||||
sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode
|
||||
sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode
|
||||
sofia_cfg.mamba_variant = cfg.sofia_mamba_variant
|
||||
sofia_cfg.mamba_backend = cfg.sofia_mamba_backend
|
||||
sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge
|
||||
sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations)
|
||||
model = SOFIAFusionEncoder(
|
||||
sofia_cfg=sofia_cfg,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
init_gate=cfg.init_gate,
|
||||
baseline_mode=cfg.baseline_mode,
|
||||
lora_rank=cfg.sofia_lora_rank,
|
||||
device=cfg.device,
|
||||
).to(cfg.device)
|
||||
elif cfg.backbone == "sofia_v1":
|
||||
sofia_v1_cfg = SOFIAv1Config(
|
||||
variant=cfg.sofia_v1_variant,
|
||||
dcn_variant=cfg.sofia_v1_dcn_variant,
|
||||
d_descriptor=cfg.sofia_v1_d_descriptor,
|
||||
text_film_dim=cfg.sofia_v1_d_descriptor,
|
||||
use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode,
|
||||
use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode,
|
||||
use_film_altitude=cfg.sofia_v1_use_film_altitude,
|
||||
)
|
||||
model = SOFIAv1FusionEncoder(
|
||||
sofia_cfg=sofia_v1_cfg,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
init_gate=cfg.init_gate,
|
||||
baseline_mode=cfg.baseline_mode,
|
||||
lora_rank=cfg.sofia_v1_lora_rank,
|
||||
device=cfg.device,
|
||||
).to(cfg.device)
|
||||
else:
|
||||
model = AsymmetricEncoder(
|
||||
dino_web_path=cfg.dino_web_path,
|
||||
dino_sat_path=cfg.dino_sat_path,
|
||||
lrsclip_path=cfg.lrsclip_path,
|
||||
init_gate=cfg.init_gate,
|
||||
baseline_mode=cfg.baseline_mode,
|
||||
shared_encoder=cfg.shared_encoder,
|
||||
mona_bottleneck=cfg.mona_bottleneck,
|
||||
mona_last_n_blocks=cfg.mona_last_n_blocks,
|
||||
device=cfg.device,
|
||||
backbone=cfg.backbone,
|
||||
stripnet_path=cfg.stripnet_path,
|
||||
stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages,
|
||||
stripnet_freeze=cfg.stripnet_freeze,
|
||||
).to(cfg.device)
|
||||
LOGGER.info("embed_dim=%d", model.embed_dim)
|
||||
|
||||
# --- Gradient checkpointing (trade compute for VRAM) ---
|
||||
# StripNet doesn't expose set_gradient_checkpointing — skip silently.
|
||||
# StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it.
|
||||
if cfg.gradient_checkpointing and cfg.backbone == "dinov3":
|
||||
if cfg.shared_encoder:
|
||||
model.image_encoder.set_gradient_checkpointing(True)
|
||||
@@ -618,10 +711,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
if model.text_encoder is not None:
|
||||
model.text_encoder.transformer.gradient_checkpointing = True
|
||||
LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
|
||||
elif cfg.gradient_checkpointing and cfg.backbone == "stripnet":
|
||||
elif cfg.gradient_checkpointing and cfg.backbone in ("stripnet", "sofia", "sofia_v1"):
|
||||
if model.text_encoder is not None:
|
||||
model.text_encoder.transformer.gradient_checkpointing = True
|
||||
LOGGER.info("Gradient checkpointing enabled (DGTRS only; StripNet doesn't support)")
|
||||
LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone)
|
||||
|
||||
n_trainable = sum(p.numel() for p in model.trainable_parameters())
|
||||
n_total = sum(p.numel() for p in model.parameters())
|
||||
@@ -879,11 +972,14 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
|
||||
drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
|
||||
sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
|
||||
altitude = batch.get("altitude")
|
||||
if altitude is not None:
|
||||
altitude = altitude.to(cfg.device, non_blocking=True)
|
||||
|
||||
# Model forward in AMP (fp16 for DINOv3/DGTRS encoders).
|
||||
with autocast(device_type="cuda", enabled=cfg.use_amp):
|
||||
if cfg.baseline_mode:
|
||||
embeddings = model(drone_img=drone_img, sat_img=sat_img)
|
||||
embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude)
|
||||
else:
|
||||
embeddings = model(
|
||||
drone_img=drone_img,
|
||||
@@ -894,6 +990,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
sat_caption_l1=batch["sat_caption_l1"],
|
||||
sat_caption_l2=batch["sat_caption_l2"],
|
||||
sat_caption_l3=batch["sat_caption_l3"],
|
||||
altitude=altitude,
|
||||
)
|
||||
# Loss — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights.
|
||||
queue_neg = neg_bank.get_queue() if neg_bank is not None else None
|
||||
@@ -1103,20 +1200,23 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
history.append(epoch_record)
|
||||
|
||||
# Save checkpoint. Model architecture flags go into the ckpt so
|
||||
# `AsymmetricEncoder.load_checkpoint` can rebuild the right shape.
|
||||
_atomic_save(
|
||||
obj={
|
||||
"epoch": epoch,
|
||||
"model_state": model.state_dict(),
|
||||
"optimizer_state": optimizer.state_dict(),
|
||||
"loss_state": loss_fn.state_dict(),
|
||||
"baseline_mode": cfg.baseline_mode,
|
||||
"shared_encoder": cfg.shared_encoder,
|
||||
"mona_bottleneck": cfg.mona_bottleneck,
|
||||
"mona_last_n_blocks": cfg.mona_last_n_blocks,
|
||||
},
|
||||
path=output_dir / f"ckpt_epoch{epoch:03d}.pt",
|
||||
)
|
||||
# `AsymmetricEncoder.load_checkpoint` (or `SOFIAFusionEncoder.load_checkpoint`)
|
||||
# can rebuild the right shape.
|
||||
ckpt_obj = {
|
||||
"epoch": epoch,
|
||||
"model_state": model.state_dict(),
|
||||
"optimizer_state": optimizer.state_dict(),
|
||||
"loss_state": loss_fn.state_dict(),
|
||||
"baseline_mode": cfg.baseline_mode,
|
||||
"backbone": cfg.backbone,
|
||||
}
|
||||
if cfg.backbone in ("sofia", "sofia_v1"):
|
||||
ckpt_obj["sofia_cfg"] = model.sofia_cfg
|
||||
else:
|
||||
ckpt_obj["shared_encoder"] = cfg.shared_encoder
|
||||
ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck
|
||||
ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks
|
||||
_atomic_save(obj=ckpt_obj, path=output_dir / f"ckpt_epoch{epoch:03d}.pt")
|
||||
LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch)
|
||||
|
||||
# Save history.
|
||||
|
||||
Reference in New Issue
Block a user