depth_edges_annotate_worlduav/docs/analysis_optimization.md

Анализ и оптимизация пайплайна 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)	Высокое	Средний	Не реализовано