# Skills: Оптимизация I/O и DL/ML инференса

Справочник приёмов оптимизации для пайплайнов обработки данных и инференса нейросетей.
Используется как чеклист при code review и проектировании новых пайплайнов.

---

## Часть 1. Оптимизация I/O

### 1.1 Асинхронная запись на диск

**Проблема:** синхронный `np.save()` / `Image.save()` блокирует основной поток — GPU простаивает.

**Решение:**

```python
from concurrent.futures import ThreadPoolExecutor

_io_pool = ThreadPoolExecutor(max_workers=4)

def save_async(fn, *args, **kwargs):
    """Отправить операцию сохранения в фоновый поток."""
    _io_pool.submit(fn, *args, **kwargs)

# В inference-цикле:
for batch in loader:
    result = model(batch)
    save_async(save_result, result, path)  # не блокирует GPU
```

**Когда применять:** всегда, когда запись идёт между батчами инференса.

**Подводные камни:**
- Контролировать размер очереди — если запись медленнее инференса, очередь растёт и съедает RAM
- При аварийном завершении незаписанные данные теряются — использовать atomic save (temp + rename)
- `_io_pool.shutdown(wait=True)` в конце пайплайна

---

### 1.2 Атомарная запись файлов

**Проблема:** прерванная запись оставляет битый файл, resume-логика считает его валидным.

**Решение:**

```python
import tempfile, os

def atomic_save_npy(arr, path):
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(suffix=".tmp", dir=path.parent)
    os.close(fd)
    try:
        np.save(tmp, arr)
        os.replace(tmp, path)   # атомарная операция на одной ФС
    except BaseException:
        os.remove(tmp) if os.path.exists(tmp) else None
        raise
```

**Когда применять:** при любом сохранении промежуточных результатов с resume-логикой.

---

### 1.3 Минимизация syscalls

**Проблема:** `Path.mkdir()`, `Path.exists()` — это syscalls. При 100k файлов × 3 стадии = 300k+ вызовов.

**Решения:**

| Приём | Описание |
|-------|----------|
| **Pre-create dirs** | Один проход `mkdir` для всех output_dir до начала обработки |
| **Completion manifest** | `set()` в памяти вместо `Path.exists()` на каждый файл |
| **Batch stat** | `os.scandir()` вместо поштучного `exists()` |

```python
# Вместо per-file exists():
completed = set()
for entry in os.scandir(output_root):
    if entry.is_dir():
        for f in os.scandir(entry.path):
            if f.name == "depth.npy":
                completed.add(entry.name)

# Фильтрация:
pending = [r for r in records if r.stem not in completed]
```

---

### 1.4 Выбор формата сохранения

| Формат | Скорость записи | Размер | Загрузка | Когда использовать |
|--------|----------------|--------|----------|--------------------|
| `.npy` (float32) | Быстро | Большой | Быстро | Промежуточные результаты |
| `.npy` (float16) | Быстро | 2× меньше | Быстро | Depth, edges [0,1] |
| `.npy` (uint8) | Быстро | 8× меньше int64 | Быстро | Class IDs (< 256 классов) |
| `.npz` compressed | Медленно | Маленький | Медленно | Архивация, не для пайплайнов |
| `.png` | Медленно (компрессия) | Маленький | Средне | Только визуализация |
| `.safetensors` | Быстро | Маленький | Быстро, mmap | Тензоры для обучения |

**Правило:** для промежуточных данных — минимальный достаточный dtype без компрессии. Визуализации — отдельный шаг по требованию.

---

### 1.5 Оптимизация PNG-записи

```python
# PIL — медленный (zlib compression level 6 по умолчанию)
Image.fromarray(vis).save("out.png")

# OpenCV — быстрее, контроль компрессии
cv2.imwrite("out.png", vis, [cv2.IMWRITE_PNG_COMPRESSION, 1])  # 1 = минимум

# Самый быстрый: отложить визуализацию
# Генерировать PNG отдельным скриптом после всего инференса
```

---

### 1.6 Кеширование файловых списков

**Проблема:** `rglob("*")` на больших датасетах — минуты.

**Решение:**

```python
import json, os

CACHE = root / ".file_cache.json"

def discover_cached(root):
    if CACHE.exists():
        mtime_cache = CACHE.stat().st_mtime
        mtime_root = root.stat().st_mtime
        if mtime_cache > mtime_root:
            return json.loads(CACHE.read_text())

    files = [str(p.relative_to(root)) for p in root.rglob("*") if p.is_file()]
    CACHE.write_text(json.dumps(files))
    return files
```

**Когда применять:** датасет > 10k файлов, особенно на HDD или сетевых ФС.

---

## Часть 2. Оптимизация DL/ML инференса

### 2.1 `torch.inference_mode()` вместо `torch.no_grad()`

```python
# Хорошо:
@torch.inference_mode()
def predict(model, x):
    return model(x)

# Менее эффективно:
with torch.no_grad():
    return model(x)
```

`inference_mode` отключает version counting и autograd tracking полностью — на 5–10% быстрее `no_grad` и меньше потребление памяти.

---

### 2.2 `torch.compile()`

```python
model = load_model()
model = torch.compile(model, mode="reduce-overhead")
# первый вызов — долгий (компиляция), далее — 20–40% быстрее
```

| Режим | Скорость компиляции | Ускорение рантайма | Когда |
|-------|--------------------|--------------------|-------|
| `"default"` | Средняя | Среднее | Общий случай |
| `"reduce-overhead"` | Медленная | Максимальное | Много батчей, стационарные формы |
| `"max-autotune"` | Очень медленная | Максимальное + подбор ядер | Продакшн, фиксированные shapes |

**Ограничения:**
- Не работает с dynamic shapes без `dynamic=True`
- Некоторые кастомные операции не компилируются — используйте `torch._dynamo.config.suppress_errors = True` при отладке
- Первый вызов медленный — прогрев обязателен

---

### 2.3 Automatic Mixed Precision (AMP)

```python
# Правильно: autocast выбирает dtype per-operation
with torch.autocast("cuda", dtype=torch.float16):
    output = model(input_tensor)

# Неправильно: ручной .half() на всё
input_tensor = input_tensor.half()
output = model(input_tensor)  # нормализация и softmax тоже в fp16 — потеря точности
```

**Что autocast делает автоматически:**
- matmul, conv → FP16 (выигрыш скорости)
- layernorm, softmax, loss → FP32 (сохранение точности)
- Accumulations → FP32

**Для обучения** дополнительно нужен `GradScaler`:

```python
scaler = torch.amp.GradScaler()
with torch.autocast("cuda", dtype=torch.float16):
    loss = model(x)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
```

---

### 2.4 Batch-инференс vs. per-image цикл

**Проблема:**

```python
# Медленно: per-image на CPU с конвертациями
for i in range(B):
    img_np = tensor_to_numpy(batch[i])     # GPU→CPU
    result = model.infer_image(img_np)      # одно изображение
    output.append(numpy_to_tensor(result))  # CPU→GPU
```

**Решение:**

```python
# Быстро: батч на GPU без конвертаций
x = preprocess(batch).to(device)  # один transfer
output = model(x)                  # весь батч за один forward
```

**Ключевые правила:**
- Избегать циклов `for i in range(B)` внутри inference-функций
- Если API модели принимает только одно изображение — обращаться к внутреннему `model.forward()` напрямую
- Минимизировать CPU↔GPU трансферы (каждый `tensor.cpu()` / `.to(device)` — задержка)
- Numpy конвертации (`tensor → numpy → PIL → numpy → tensor`) — красный флаг

---

### 2.5 CUDA Streams для overlap

```python
stream_compute = torch.cuda.Stream()
stream_transfer = torch.cuda.Stream()

for batch in loader:
    with torch.cuda.stream(stream_transfer):
        next_batch = next_batch.to(device, non_blocking=True)

    with torch.cuda.stream(stream_compute):
        output = model(current_batch)

    torch.cuda.synchronize()
```

**Когда применять:** per-image модели (SegEarth/SAM), где нельзя батчить — overlap загрузки следующего изображения с инференсом текущего.

---

### 2.6 DataLoader — скрытые параметры производительности

```python
DataLoader(
    dataset,
    batch_size=bs,
    num_workers=4,              # параллельная загрузка
    pin_memory=True,            # ускоряет CPU→GPU transfer
    persistent_workers=True,    # не пересоздавать процессы между эпохами
    prefetch_factor=4,          # буфер загрузки (дефолт 2)
    drop_last=False,            # True для обучения, False для инференса
)
```

**Диагностика bottleneck загрузки:**

```python
import time
for batch in loader:
    t0 = time.perf_counter()
    output = model(batch)
    gpu_time = time.perf_counter() - t0

    t0 = time.perf_counter()
    # просто итерация — если это медленно, bottleneck в загрузке
load_time = time.perf_counter() - t0
```

Если `load_time > gpu_time` — увеличить `num_workers`, добавить `prefetch_factor`, перейти на DALI.

---

### 2.7 Decode и preprocessing на GPU

```python
# CPU (медленно для больших датасетов):
img = Image.open(path).convert("RGB")
tensor = transforms.ToTensor()(transforms.Resize((H, W))(img))

# GPU через torchvision v2:
from torchvision.io import decode_image
from torchvision.transforms import v2

tensor = decode_image(path)  # decode на CPU
tensor = tensor.to(device)
tensor = v2.Resize((H, W))(tensor)  # resize на GPU

# NVIDIA DALI (полностью на GPU):
# decode JPEG → resize → normalize → выход как CUDA тензор
# Требует отдельного pipeline definition, но 3–5× быстрее для >10k images
```

---

### 2.8 Оптимизация размера батча

```python
def auto_batch_size(total_vram_mb, weights_mb, per_sample_mb, overhead_mb=200):
    """Вычислить безопасный batch size по VRAM бюджету."""
    free = total_vram_mb - weights_mb - overhead_mb
    if free <= 0:
        return 1
    max_b = int(free / per_sample_mb)
    # Округление вниз до степени 2 (лучше для GPU)
    safe = max(1, 2 ** int(math.log2(max(1, int(max_b * 0.7)))))
    return safe
```

**Правила:**
- Степени двойки — лучшая утилизация GPU tensor cores
- 70% от теоретического максимума — запас на пиковое потребление
- `torch.cuda.mem_get_info()` — актуальная свободная память в рантайме (лучше статических оценок)
- При OOM — fallback на `batch_size // 2` с retry

---

### 2.9 Оптимизация конвертаций данных

**Красные флаги (каждый — потеря производительности):**

| Паттерн | Проблема | Решение |
|---------|----------|---------|
| `tensor.numpy()` в цикле | GPU→CPU sync | Батчить, делать `.cpu()` один раз |
| `np.uint8 → float32 → device` | Двойная конвертация | Конвертировать сразу на device |
| `Image.fromarray()` в цикле | PIL overhead | cv2 или отложить |
| `tensor.item()` в цикле | Sync point | Собирать в тензор, `.item()` один раз |
| Пересоздание тензоров-констант | Аллокация памяти | Module-level или `register_buffer` |

---

### 2.10 Профилирование инференса

```python
# Быстрая диагностика: где время?
import torch.utils.benchmark as benchmark

timer = benchmark.Timer(
    stmt="model(x)",
    globals={"model": model, "x": sample_batch},
    num_threads=1,
)
print(timer.blocked_autorange())

# Детальный профиль:
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,
    with_stack=True,
) as prof:
    model(sample_batch)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
# Экспорт в Chrome trace:
prof.export_chrome_trace("trace.json")
```

**Что искать:**
- Операции с большим `cuda_time` — кандидаты для `torch.compile` / fusion
- Большое `cpu_time` при маленьком `cuda_time` — Python overhead, data transfer
- Много мелких CUDA kernels — нужен fusion (compile / scripting)

---

### 2.11 Кеширование повторяющихся тензоров

```python
# Плохо: пересоздание при каждом вызове
def process(x):
    kernel = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32)
    return F.conv2d(x, kernel.view(1, 1, 3, 3))

# Хорошо: module-level константа
_SOBEL_X = torch.tensor(
    [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32
).view(1, 1, 3, 3) / 8.0

def process(x):
    kernel = _SOBEL_X.to(x.device)  # перенос на device только при смене
    return F.conv2d(x, kernel)
```

---

### 2.12 Экспорт и рантаймы для продакшн-инференса

| Рантайм | Ускорение | Сложность | Когда |
|---------|-----------|-----------|-------|
| `torch.compile` | 1.2–1.5× | Низкая | Первый шаг оптимизации |
| TorchScript (`torch.jit.trace`) | 1.1–1.3× | Средняя | Деплой без Python |
| ONNX Runtime | 1.3–2× | Средняя | Кросс-платформенный инференс |
| TensorRT | 2–5× | Высокая | NVIDIA GPU, фиксированные shapes |
| OpenVINO | 1.5–3× | Средняя | Intel CPU/GPU |

```python
# ONNX экспорт:
torch.onnx.export(model, sample_input, "model.onnx",
                  input_names=["image"], output_names=["depth"],
                  dynamic_axes={"image": {0: "batch"}, "depth": {0: "batch"}})

# TensorRT через torch_tensorrt:
import torch_tensorrt
trt_model = torch_tensorrt.compile(model,
    inputs=[torch_tensorrt.Input(shape=[bs, 3, H, W], dtype=torch.float16)],
    enabled_precisions={torch.float16})
```

---

## Часть 3. Чеклист для нового пайплайна

### Перед написанием кода

- [ ] Определить bottleneck: GPU compute, CPU preprocessing, или disk I/O?
- [ ] Выбрать минимально достаточные dtype для хранения
- [ ] Спланировать overlap: загрузка ↔ инференс ↔ сохранение

### Инференс

- [ ] `torch.inference_mode()` на всех predict-функциях
- [ ] `torch.compile()` если > 100 батчей
- [ ] `torch.autocast` вместо ручного `.half()`
- [ ] Batch forward вместо per-image циклов
- [ ] Нет лишних CPU↔GPU трансферов
- [ ] Нет numpy конвертаций внутри inference loop
- [ ] Тензоры-константы вынесены на module level

### Data Pipeline

- [ ] `pin_memory=True`
- [ ] `persistent_workers=True`
- [ ] `prefetch_factor` ≥ 2
- [ ] `num_workers` подобран (обычно 4–8)
- [ ] Decode/resize на GPU при > 10k images

### I/O

- [ ] Асинхронная запись через `ThreadPoolExecutor`
- [ ] Atomic save для crash safety
- [ ] Минимум syscalls (pre-create dirs, batch stat)
- [ ] Визуализации отделены от основного пайплайна
- [ ] Completion tracking через manifest, не per-file `exists()`

### Память

- [ ] `gc.collect()` + `torch.cuda.empty_cache()` между стадиями
- [ ] Модели выгружаются когда не нужны
- [ ] Промежуточные тензоры не удерживаются в памяти
- [ ] dtype соответствует данным (uint8 для class IDs, float16 для [0,1])
