Initial commit: World-UAV annotation pipeline
4-modality annotation pipeline (depth, edges, segmentation, chmv2) for 973K drone/satellite images. SegEarth-OV3 open-vocabulary segmentation with 11 classes optimized for cross-view geo-localization. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
326
docs/segearth_ov3_architecture.md
Normal file
326
docs/segearth_ov3_architecture.md
Normal file
@@ -0,0 +1,326 @@
|
||||
# SegEarth-OV3: архитектура, pipeline и оптимизация
|
||||
|
||||
## Обзор
|
||||
|
||||
SegEarth-OV3 — модель **open-vocabulary** семантической сегментации для дистанционного зондирования, построенная на базе **SAM 3.1** (Segment Anything Model 3.1, Meta). Позволяет сегментировать изображения по произвольным текстовым описаниям классов без дообучения.
|
||||
|
||||
В нашем пайплайне используется для генерации семантических карт из аэрофотоснимков (дрон) и спутниковых изображений датасета World-UAV (973K изображений).
|
||||
|
||||
---
|
||||
|
||||
## Архитектура модели
|
||||
|
||||
### Общая схема
|
||||
|
||||
```
|
||||
┌──────────────────────────────────────────┐
|
||||
│ SAM3VLBackbone │
|
||||
│ ┌─────────────────┐ ┌───────────────┐ │
|
||||
Image ──────────►│ │ Vision Backbone │ │ Text Encoder │◄─── Text Prompt
|
||||
[B,3,1008,1008] │ │ (ViT + Neck) │ │ (VE, 24-layer)│ │
|
||||
│ └────────┬────────┘ └───────┬───────┘ │
|
||||
│ │ │ │
|
||||
│ vision_features language_features│
|
||||
│ [B,256,72,72] [32,1,256] │
|
||||
└──────────┬────────────────────┬───────────┘
|
||||
│ │
|
||||
▼ ▼
|
||||
┌──────────────────────────────────────────┐
|
||||
│ Transformer Encoder (6 layers) │
|
||||
│ Cross-attention: vision × language │
|
||||
│ Self-attention: vision features │
|
||||
└──────────────────┬───────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────────────────────────────┐
|
||||
│ Transformer Decoder (6 layers) │
|
||||
│ 200 object queries │
|
||||
│ Cross-attention: queries × encoder out │
|
||||
│ Text cross-attention │
|
||||
│ Box refinement (DAC) │
|
||||
│ Presence token │
|
||||
└──────────┬───────────────────────────────┘
|
||||
│
|
||||
┌──────────┴──────────┐
|
||||
│ │
|
||||
▼ ▼
|
||||
┌──────────────┐ ┌───────────────────┐
|
||||
│ Instance Head │ │ Segmentation Head │
|
||||
│ pred_boxes │ │ (PixelDecoder + │
|
||||
│ pred_masks │ │ semantic_seg) │
|
||||
│ pred_logits │ │ │
|
||||
│ presence_score│ │ semantic_mask_logits│
|
||||
└──────────────┘ └───────────────────┘
|
||||
│ │
|
||||
└──────────┬──────────┘
|
||||
▼
|
||||
┌──────────────────────────────────────────┐
|
||||
│ Post-processing │
|
||||
│ argmax → class index map │
|
||||
│ threshold filtering → background │
|
||||
└──────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Компоненты
|
||||
|
||||
#### 1. Vision Backbone (ViT + FPN Neck)
|
||||
|
||||
- **ViT** (Vision Transformer): 32 слоя, embed_dim=1024, 16 heads
|
||||
- Входное разрешение: **1008×1008** (patch_size=14 → 72×72 patches)
|
||||
- RoPE (Rotary Position Embeddings) с интерполяцией
|
||||
- Window attention (window=24) + 4 global attention blocks (слои 7, 15, 23, 31)
|
||||
- Tiled absolute position embeddings (pretrain=336 → tile до 1008)
|
||||
- **Sam3DualViTDetNeck**: FPN-neck с 4 масштабами (×4, ×2, ×1, ×0.5)
|
||||
- Выход: `[B, 256, 72, 72]` (после scalp=1, отбрасывается низшее разрешение)
|
||||
- Также генерирует SAM2-совместимые features для instance interactivity
|
||||
|
||||
#### 2. Text Encoder (VETextEncoder)
|
||||
|
||||
- 24-слойный Transformer, width=1024, 16 heads
|
||||
- BPE токенизатор (vocab: 16M tokens из `bpe_simple_vocab_16e6.txt.gz`)
|
||||
- Вход: текстовый промпт (например, `"railway"`)
|
||||
- Выход:
|
||||
- `language_features`: `[32, 1, 256]` — проекция в 256-dim space
|
||||
- `language_mask`: `[1, 32]` — маска внимания
|
||||
- `language_embeds`: `[32, 1, 1024]` — полные эмбеддинги
|
||||
|
||||
#### 3. Transformer Encoder (6 layers)
|
||||
|
||||
- Fusion: cross-attention между vision features и language features
|
||||
- Self-attention на vision features
|
||||
- d_model=256, dim_feedforward=2048, 8 heads
|
||||
- Activation checkpointing включён
|
||||
|
||||
#### 4. Transformer Decoder (6 layers)
|
||||
|
||||
- **200 object queries** — learnable queries для обнаружения объектов
|
||||
- Cross-attention: queries × encoder output
|
||||
- Text cross-attention: queries × language features
|
||||
- **DAC** (Decoupled Attention for Classification) — разделение box regression и classification
|
||||
- **Box refinement** на каждом слое (iterative)
|
||||
- **Presence token** — для оценки наличия объекта в сцене
|
||||
- Resolution=1008, stride=14
|
||||
|
||||
#### 5. Scoring & Heads
|
||||
|
||||
- **DotProductScoring**: MLP (256→2048→256) + dot product для class scores
|
||||
- **Instance Head**: pred_boxes, pred_masks, pred_logits, presence_score
|
||||
- **Segmentation Head** (UniversalSegmentationHead):
|
||||
- PixelDecoder: 3-stage upsampling (nearest interpolation), hidden_dim=256
|
||||
- Выход: `semantic_seg` — пиксельная семантическая маска
|
||||
|
||||
---
|
||||
|
||||
## Pipeline инференса
|
||||
|
||||
### Для одного изображения (`predict_pil`)
|
||||
|
||||
```
|
||||
1. set_image(pil_image)
|
||||
└── transform → [1, 3, 1008, 1008]
|
||||
└── backbone.forward_image() → backbone_out (vision features)
|
||||
|
||||
2. Для каждого из 11 промптов:
|
||||
a. reset_all_prompts(state)
|
||||
b. set_text_prompt(prompt, state)
|
||||
└── backbone.forward_text([prompt]) → text_outputs
|
||||
└── state["backbone_out"].update(text_outputs)
|
||||
└── _forward_grounding(state)
|
||||
├── model.forward_grounding(backbone_out, geometric_prompt)
|
||||
├── pred_boxes, pred_masks, pred_logits → filter by confidence
|
||||
└── semantic_seg → interpolate to (H, W)
|
||||
c. Агрегация:
|
||||
- instance: seg_logits[i] = max(seg_logits[i], mask * score)
|
||||
- semantic: seg_logits[i] = max(seg_logits[i], semantic_logits)
|
||||
- presence: seg_logits[i] *= presence_score
|
||||
|
||||
3. Post-processing:
|
||||
└── argmax(seg_logits, dim=0) → class map
|
||||
└── max_vals < threshold → set to background (class 0)
|
||||
```
|
||||
|
||||
**Итого:** 1 backbone pass + **11 × (text_encoder + grounding_decoder)** forward passes.
|
||||
|
||||
### Для батча (`predict_pil_batch`)
|
||||
|
||||
```
|
||||
1. set_image_batch(images)
|
||||
└── transform all → [B, 3, 1008, 1008]
|
||||
└── backbone.forward_image(batch) → batch_state ← ОДНА операция на весь батч
|
||||
|
||||
2. Для каждого изображения i в батче:
|
||||
a. _slice_backbone_out(batch_state, i) → per-image state
|
||||
b. Для каждого из 11 промптов:
|
||||
└── (используются кэшированные text embeddings)
|
||||
└── _forward_grounding(state) → instance + semantic masks
|
||||
c. argmax + threshold → class map
|
||||
```
|
||||
|
||||
**Батчинг backbone** экономит ~15% времени vs per-image.
|
||||
**Кэширование text embeddings** экономит ~2% (text encoder очень быстрый).
|
||||
|
||||
---
|
||||
|
||||
## Профиль производительности (RTX 4090, 24 GB VRAM)
|
||||
|
||||
### Разбивка по времени (на 1 изображение, 256×256, 11 промптов)
|
||||
|
||||
| Этап | Время | Доля |
|
||||
|---|---|---|
|
||||
| Vision backbone (ViT + Neck) | ~25 ms | ~9% |
|
||||
| Text encoder (× 11 промптов) | ~5 ms | ~2% |
|
||||
| Grounding decoder (× 11 промптов) | ~240 ms | ~84% |
|
||||
| Post-processing (argmax, threshold) | ~2 ms | ~1% |
|
||||
| Overhead (PIL convert, transfer) | ~14 ms | ~5% |
|
||||
| **Итого** | **~286 ms** | **100%** |
|
||||
|
||||
### Throughput при разных batch size
|
||||
|
||||
| Batch size | Throughput | VRAM | ms/img |
|
||||
|---|---|---|---|
|
||||
| 8 | 3.4 img/s | 6.2 GB | 292 |
|
||||
| 10 | 3.5 img/s | 6.2 GB | 286 |
|
||||
| 16 | 3.3 img/s | 8.1 GB | 301 |
|
||||
| 24 | 3.5 img/s | 8.8 GB | 286 |
|
||||
| 32 | 3.5 img/s | 8.8 GB | 284 |
|
||||
|
||||
**Вывод:** throughput **не масштабируется** с batch size, т.к. bottleneck — **per-image grounding decoder** (11 последовательных forward passes на каждое изображение). Backbone батчится, но это лишь 9% времени.
|
||||
|
||||
### Оценка на полный датасет
|
||||
|
||||
| Подмножество | Кол-во изображений | Время при 3.5 img/s |
|
||||
|---|---|---|
|
||||
| DB (спутник) | ~486K | ~38.5 часов |
|
||||
| Query (дрон) | ~486K | ~38.5 часов |
|
||||
| Всё | ~973K | ~77 часов |
|
||||
|
||||
---
|
||||
|
||||
## Применённые оптимизации
|
||||
|
||||
### 1. Кэширование text embeddings
|
||||
|
||||
**Файл:** `src/nn/segearth_ov3/segearthov3_segmentor.py`
|
||||
|
||||
Text encoder вызывается с одними и теми же 11 промптами для каждого изображения. Кэширование результатов `forward_text()` при первом вызове и повторное использование для всех последующих изображений.
|
||||
|
||||
```python
|
||||
def _cache_text_embeddings(self):
|
||||
"""Pre-compute and cache text embeddings for all prompts (run once)."""
|
||||
if hasattr(self, '_text_cache'):
|
||||
return
|
||||
self._text_cache = []
|
||||
for query_word in self.query_words:
|
||||
text_out = self.processor.model.backbone.forward_text(
|
||||
[query_word], device=self.device,
|
||||
)
|
||||
cached = {k: v.clone() if isinstance(v, torch.Tensor) else v
|
||||
for k, v in text_out.items()}
|
||||
self._text_cache.append(cached)
|
||||
```
|
||||
|
||||
**Эффект:** ~2% ускорение (54 ms на батч из 8). Text encoder итак быстрый (~0.5 ms/промпт).
|
||||
|
||||
### 2. Увеличение batch size (8 → 16)
|
||||
|
||||
**Файлы:** `src/main.py`, `src/augmentor/inference.py`
|
||||
|
||||
- `_MAX_SEG_BATCH`: 8 → 16
|
||||
- Segmentation stage `bs`: 8 → 16
|
||||
|
||||
**Эффект:** VRAM вырос 6.2 → 8.8 GB (из 24 доступных). Throughput стабилен, но меньше overhead на создание батчей и DataLoader.
|
||||
|
||||
### 3. Autocast bfloat16
|
||||
|
||||
Уже включён в оригинальном коде:
|
||||
```python
|
||||
with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
||||
```
|
||||
|
||||
### 4. TF32 для Ampere+ GPU
|
||||
|
||||
Включён автоматически в `model_builder.py`:
|
||||
```python
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Фундаментальные ограничения (нельзя ускорить без переписывания модели)
|
||||
|
||||
### 1. Последовательный grounding per prompt
|
||||
|
||||
Grounding decoder принимает **один текстовый промпт за раз** и прогоняет весь decoder (6 layers, 200 queries). При 11 промптах это 11 × decoder forward = **84% общего времени**.
|
||||
|
||||
**Почему нельзя батчить:** архитектура decoder требует `language_features` в cross-attention — разные промпты дают разные `language_features`, что меняет attention pattern. Для батчинга нужно переписать decoder для поддержки multi-prompt inference (нетривиально).
|
||||
|
||||
### 2. Последовательный grounding per image
|
||||
|
||||
Даже при батчинге backbone, grounding всё равно выполняется **per-image** (slice backbone features → per-image state → 11 decoder passes). Причина: decoder выдаёт per-image instance predictions (boxes, masks, scores), которые нельзя батчить из-за переменного числа detected instances.
|
||||
|
||||
### 3. Высокое разрешение ViT
|
||||
|
||||
ViT работает на 1008×1008 (72×72 patches), что требует значительного compute. Снижение `default_resolution` ускорит backbone, но ухудшит качество сегментации мелких объектов.
|
||||
|
||||
---
|
||||
|
||||
## Альтернативные пути ускорения (не реализованы)
|
||||
|
||||
| Подход | Ожидаемый эффект | Сложность | Риск |
|
||||
|---|---|---|---|
|
||||
| `torch.compile(model)` | 10-30% на decoder | Средняя | Dynamic shapes могут сломать |
|
||||
| Снижение `num_queries` (200 → 100) | ~20% на decoder | Нужно переучить | Потеря мелких объектов |
|
||||
| Снижение `default_resolution` (1008 → 504) | ~4x backbone | Тривиально (config) | Ухудшение качества |
|
||||
| Multi-GPU inference | ~2x при 2 GPU | Средняя | Нужен второй GPU |
|
||||
| ONNX/TensorRT export | 2-5x overall | Высокая | SAM3 dynamic shapes |
|
||||
| Замена на SegFormer | ~3x быстрее | Тривиально (fallback есть) | Нет open-vocab, фиксированные 150 классов |
|
||||
|
||||
---
|
||||
|
||||
## Конфигурация в проекте
|
||||
|
||||
### Файлы
|
||||
|
||||
| Файл | Назначение |
|
||||
|---|---|
|
||||
| `in/config_files/segmentation.gin` | Промпты, threshold, разрешение |
|
||||
| `src/nn/segearth_ov3/segearthov3_segmentor.py` | Обёртка SegEarth-OV3 (predict_pil, predict_pil_batch) |
|
||||
| `src/nn/segearth_ov3/sam3/model_builder.py` | Сборка модели (build_sam3_image_model) |
|
||||
| `src/nn/segearth_ov3/sam3/model/sam3_image_processor.py` | Inference processor (set_image, set_text_prompt) |
|
||||
| `src/nn/segearth_ov3/sam3/model/sam3_image.py` | Основная модель Sam3Image |
|
||||
| `src/nn/segearth_ov3/sam3/model/vl_combiner.py` | Vision-Language backbone |
|
||||
| `src/augmentor/models.py` | Загрузка модели (load_segmentation_model) |
|
||||
| `src/augmentor/inference.py` | Батчевый инференс (infer_segmentation_batch) |
|
||||
| `in/weights/sam3.1/sam3.1_multiplex.pt` | Чекпоинт модели (~2.5 GB) |
|
||||
|
||||
### Текущая конфигурация
|
||||
|
||||
```gin
|
||||
SegConfig.prompts = [
|
||||
'background', # 0
|
||||
'building', # 1
|
||||
'road', # 2
|
||||
'vegetation', # 3
|
||||
'water', # 4
|
||||
'sand and gravel ground', # 5
|
||||
'rocky terrain', # 6
|
||||
'farmland', # 7
|
||||
'railway', # 8
|
||||
'parking lot', # 9
|
||||
'sidewalk', # 10
|
||||
]
|
||||
SegConfig.threshold = 0.15
|
||||
SegConfig.default_resolution = 1008
|
||||
```
|
||||
|
||||
### Параметры модели
|
||||
|
||||
| Параметр | Значение | Описание |
|
||||
|---|---|---|
|
||||
| `prob_thd` | 0.15 | Минимальная уверенность для не-background класса |
|
||||
| `confidence_threshold` | 0.5 | Порог для instance detection (boxes/masks) |
|
||||
| `use_sem_seg` | True | Использовать semantic segmentation head |
|
||||
| `use_presence_score` | True | Масштабировать logits на presence score |
|
||||
| `use_transformer_decoder` | True | Использовать instance-level predictions |
|
||||
| `bg_idx` | 0 | Индекс класса background |
|
||||
Reference in New Issue
Block a user