depth_edges_annotate_worlduav/docs/segearth_ov3_architecture.md

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() при первом вызове и повторное использование для всех последующих изображений.

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

Уже включён в оригинальном коде:

with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):

4. TF32 для Ampere+ GPU

Включён автоматически в model_builder.py:

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)

Текущая конфигурация

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