686db62c25
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>
164 lines
9.9 KiB
Markdown
164 lines
9.9 KiB
Markdown
# Анализ и оптимизация пайплайна depth_edges_annotate_worlduav
|
||
|
||
## Обзор
|
||
|
||
Пайплайн для аннотации датасета World-UAV: четыре стадии (depth, edges, segmentation, chmv2), конфигурация через Gin, fallback-модели, атомарное сохранение, resume-логика. Нейросетевые модели (SegEarth-OV-3, Depth-Anything-3) вендорированы в `src/nn/`.
|
||
|
||
---
|
||
|
||
## 1. Инференс GPU -- самые крупные выигрыши
|
||
|
||
### 1.1 `torch.compile()` для моделей
|
||
|
||
В `models.py` после загрузки модели можно обернуть в `torch.compile(model, mode="reduce-overhead")`. Для DepthAnything и SegFormer это может дать **20-40% ускорения** на repeated batches за счет fusion ядер и устранения Python overhead. Особенно эффективно при большом количестве батчей.
|
||
|
||
### 1.2 AMP через `torch.autocast` вместо ручного `.half()`
|
||
|
||
Сейчас в `inference.py:79` и `inference.py:203` делается ручной каст `x = x.half()` при условии `model.parameters().dtype == torch.float16`. Это неоптимально:
|
||
|
||
- `torch.autocast("cuda", dtype=torch.float16)` автоматически выберет FP16 для matmul/conv, но оставит FP32 для нормализации и softmax -- меньше потерь точности и обычно быстрее
|
||
- Убирает необходимость в ручных проверках dtype
|
||
|
||
### 1.3 DA3: batch-инференс вместо поштучного цикла
|
||
|
||
`inference.py:59-74` -- DA3 обрабатывается через `model.inference()` API, который внутри делает preprocess + forward + postprocess. Для максимальной производительности можно вызывать низкоуровневый `model.forward()` напрямую (весь батч тензором на GPU), минуя конвертации tensor -> numpy -> PIL -> обратно.
|
||
|
||
### 1.4 SegEarth-OV3: batched backbone (реализовано в v3.2)
|
||
|
||
`predict_pil_batch()` батчит backbone SAM 3.1 (~80% времени) на до 8 изображений. Grounding decoder остается per-image. Дальнейшие оптимизации:
|
||
|
||
- `torch.cuda.Stream` для overlap compute/data-transfer
|
||
- Увеличение max_batch если VRAM позволяет (текущий лимит 8 hardcoded)
|
||
|
||
---
|
||
|
||
## 2. Data Pipeline -- ускорение загрузки
|
||
|
||
### 2.1 `persistent_workers=True` (реализовано)
|
||
|
||
DataLoader создается с `persistent_workers=True` при `num_workers > 0`. Воркеры не пересоздаются между батчами.
|
||
|
||
### 2.2 `prefetch_factor` (реализовано)
|
||
|
||
`prefetch_factor=4` для лучшего overlap загрузки с инференсом.
|
||
|
||
### 2.3 Decode на GPU через `torchvision.io` или `nvidia-dali`
|
||
|
||
`dataset.py` -- `Image.open().convert("RGB")` + `transforms.Resize()` работает на CPU через PIL. Для больших датасетов это bottleneck. Варианты:
|
||
|
||
- `torchvision.io.decode_image` + `torchvision.transforms.v2` -- resize на GPU
|
||
- NVIDIA DALI pipeline -- полный decode+resize+normalize на GPU, особенно выгоден при images > 10k
|
||
|
||
### 2.4 Предвычисление списка файлов
|
||
|
||
`dataset.py` -- `rglob("*")` обходит файловую систему при каждом запуске. Для больших датасетов (тысячи папок) это минуты. Можно кешировать список в `.file_cache.json` и обновлять только при изменении mtime корня.
|
||
|
||
---
|
||
|
||
## 3. I/O -- запись результатов
|
||
|
||
### 3.1 Асинхронная запись (реализовано)
|
||
|
||
`io_utils.py` -- `ThreadPoolExecutor` с 4 workers для неблокирующей записи файлов. GPU inference продолжается пока предыдущий батч пишется на диск.
|
||
|
||
### 3.2 Атомарная запись (реализовано)
|
||
|
||
Temp file + `os.replace()` для crash-safety. Resume-логика корректно обрабатывает прерванные записи.
|
||
|
||
### 3.3 Визуализации -- отложить или отключить
|
||
|
||
`save_vis=False` отключает PNG-визуализации. Для ускорения можно генерировать визуализации отдельным скриптом после всего инференса.
|
||
|
||
### 3.4 Pre-create output dirs (реализовано)
|
||
|
||
Все output_dir создаются одним проходом до начала обработки (`run_pipeline()` в `main.py:333-339`), а не per-image.
|
||
|
||
---
|
||
|
||
## 4. Edges стадия
|
||
|
||
### 4.1 Batched Sobel (реализовано)
|
||
|
||
Edges обрабатываются батчами по 32 изображения. `compute_edges_from_depth` поддерживает `[B, 1, H, W]`.
|
||
|
||
### 4.2 Sobel-ядра как module-level константы (реализовано)
|
||
|
||
`_SOBEL_X` и `_SOBEL_Y` -- module-level константы в `inference.py`.
|
||
|
||
### 4.3 Edges как побочный продукт depth
|
||
|
||
Если не нужен resume между стадиями -- считать edges прямо в `run_depth_stage` из только что вычисленного depth, не сохраняя и не загружая `.npy`. Текущая архитектура предпочитает гранулярный resume.
|
||
|
||
---
|
||
|
||
## 5. Resume и discovery
|
||
|
||
### 5.1 Completion manifest вместо per-file проверок
|
||
|
||
`filter_completed()` -- для каждого изображения вызывается `Path.exists()` (syscall). При 100k изображений x 4 стадии = 400k stat-вызовов. Альтернатива:
|
||
|
||
- `completed.json` / SQLite per stage
|
||
- Проверка по set-lookup вместо filesystem
|
||
|
||
### 5.2 `filter_completed` -- per-stage
|
||
|
||
Каждая стадия фильтрует отдельно. Можно за один проход собрать статусы всех стадий.
|
||
|
||
---
|
||
|
||
## 6. Память и типы данных
|
||
|
||
### 6.1 FP32 для CHMv2 (вынужденная мера)
|
||
|
||
CHMv2 (DINOv3 DPT head) выдает NaN в FP16. Всегда FP32 -- это увеличивает VRAM на ~0.65 GB vs FP16, но гарантирует корректность.
|
||
|
||
### 6.2 uint8 для сегментации (реализовано)
|
||
|
||
`infer_segmentation_batch` возвращает `uint8` для class IDs. При 5 классах это оптимально.
|
||
|
||
### 6.3 float16 для .npy (реализовано)
|
||
|
||
Depth, edges, CHM сохраняются в `.npy` как `float16` (если `save_npy=True`).
|
||
|
||
---
|
||
|
||
## 7. Вендорированные пакеты (src/nn/)
|
||
|
||
### 7.1 Текущая архитектура
|
||
|
||
SegEarth-OV-3 и Depth-Anything-3 **встроены в проект** как вендорированные пакеты в `src/nn/`. При `import src.nn` автоматически регистрируются пути в `sys.path`:
|
||
- `src/nn/` -- для `depth_anything_3.*`
|
||
- `src/nn/segearth_ov3/` -- для `sam3.*` и `segearthov3_segmentor`
|
||
|
||
**Преимущества:**
|
||
- Нет зависимости от внешних репозиториев
|
||
- Воспроизводимость -- фиксированная версия кода
|
||
- Нет конфликтов с системными пакетами
|
||
|
||
**Ограничения:**
|
||
- Обновление до новой версии модели требует ручного копирования
|
||
- Дублирование кода (~6.5 MB) -- допустимо для проекта с ~200 GB данных
|
||
|
||
### 7.2 BPE vocab
|
||
|
||
Файл `bpe_simple_vocab_16e6.txt.gz` встроен в `src/nn/segearth_ov3/sam3/assets/`. Дефолтный путь в `segearthov3_segmentor.py` вычисляется через `os.path.dirname(__file__)`. Копия в `in/weights/` используется приоритетно (если существует) для совместимости с существующими инсталляциями.
|
||
|
||
---
|
||
|
||
## Приоритет по соотношению усилие/выигрыш
|
||
|
||
| # | Оптимизация | Усилие | Выигрыш | Статус |
|
||
|---|-------------|--------|---------|--------|
|
||
| 1 | Async I/O через ThreadPool | Низкое | **Высокий** (overlap GPU/disk) | **Реализовано** |
|
||
| 2 | Batched backbone SegEarth-OV3 | Среднее | **Высокий** (~2.5-3x segm) | **Реализовано** |
|
||
| 3 | Pre-create output dirs | Минимальное | Средний | **Реализовано** |
|
||
| 4 | Batched Sobel edges | Низкое | Средний | **Реализовано** |
|
||
| 5 | dtype: uint8 для seg, float16 для npy | Низкое | Средний | **Реализовано** |
|
||
| 6 | `persistent_workers` + `prefetch_factor` | Минимальное | Низкий-Средний | **Реализовано** |
|
||
| 7 | Вендоринг моделей (src/nn/) | Среднее | Средний (reliability) | **Реализовано** |
|
||
| 8 | `torch.compile()` | Низкое | Средний (20-40%) | Не реализовано |
|
||
| 9 | `torch.autocast` вместо `.half()` | Низкое | Низкий-Средний | Не реализовано |
|
||
| 10 | DA3 прямой forward вместо `inference()` | Среднее | Средний | Не реализовано |
|
||
| 11 | Completion manifest вместо per-file exists | Среднее | Средний (при >100k img) | Не реализовано |
|
||
| 12 | GPU decode (DALI / torchvision.io) | Высокое | Средний | Не реализовано |
|