Какие знаешь методы ускорения нейросетей?

Методы ускорения нейросетей

Ускорение нейросетей — критическая задача для развёртывания моделей в production и достижения приемлемых времён предсказания. Существует множество подходов на разных уровнях абстракции.

1. Квантизация (Quantization)

Описание — уменьшение точности представления весов и активаций (с float32 до int8, int16).

INT8 квантизация — наиболее практичная:

import torch
from torch.quantization import quantize_dynamic

model_quantized = quantize_dynamic(
    model, 
    {torch.nn.Linear}, 
    dtype=torch.qint8
)

Post-Training Quantization (PTQ) — квантизация уже обученной модели, не требует переобучения.

Quantization-Aware Training (QAT) — учитывает квантизацию во время обучения, даёт лучшую точность:

from torch.quantization import QConfig, prepare_qat, convert

model.qconfig = torch.quantization.get_default_qat_qconfig(fbgemm)
prepare_qat(model, inplace=True)
# обучение...
convert(model, inplace=True)

Результат: 2-4x ускорение, снижение памяти на 75%.

2. Pruning (Обрезка)

Удаление малозначимых весов — близких к нулю или с низкой чувствительностью:

import torch.nn.utils.prune as prune

prune.l1_unstructured(layer, name=weight, amount=0.3)  # удалить 30%
prune.remove(layer, weight)  # сделать постоянным

Structured pruning — удаляет целые каналы/фильтры, получается ускорение на железе:

# Удаление каналов
prune.ln_structured(layer, name=weight, amount=0.2, n=2, dim=0)

Результат: 5-10x ускорение при минимальной потере точности.

3. Дистилляция (Knowledge Distillation)

Идея — обучить маленькую student-модель воспроизводить поведение большой teacher-модели:

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, targets, T=4, alpha=0.7):
    ce_loss = F.cross_entropy(student_logits, targets)
    kl_loss = F.kl_div(
        F.log_softmax(student_logits / T, dim=1),
        F.softmax(teacher_logits / T, dim=1),
        reduction=batchmean
    ) * (T ** 2)
    return alpha * ce_loss + (1 - alpha) * kl_loss

Результат: 10-50x ускорение в inference, часто с минимальной потерей точности.

4. Оптимизация архитектуры

MobileNet/EfficientNet — специально разработаны для мобильных устройств:

• Depthwise separable convolutions вместо обычных
• Меньше параметров, аналогичная точность

Pruned архитектуры — AutoML инструменты ищут оптимальные архитектуры (NAS):

# Использование готовой облегчённой версии
model = torch.hub.load(pytorch/vision:v0.10.0, mobilenet_v2, pretrained=True)

5. Оптимизация на уровне железа

Batch normalization folding — объединить BN слой с предыдущим convolution:

from torch.nn.utils.fusion import fuse_conv_bn

fused_model = fuse_conv_bn(model)

Использование GPU/TPU — параллельное вычисление вместо CPU.

TensorRT (NVIDIA) — компилятор для оптимизации моделей на NVIDIA GPU:

trtexec --onnx=model.onnx --saveEngine=model.trt

ONNX Runtime — кроссплатформенная оптимизация:

import onnxruntime as ort

sess = ort.InferenceSession(model.onnx)
outputs = sess.run(None, {input: input_data})

6. Low-rank разложение

SVD (Singular Value Decomposition) — разложить матрицу весов на произведение меньших матриц:

U, S, Vh = np.linalg.svd(weight_matrix, full_matrices=False)
rank = 50  # оставить только top-50 сингулярных значений
weight_reduced = U[:, :rank] @ np.diag(S[:rank]) @ Vh[:rank, :]

LoRA (Low-Rank Adaptation) — для fine-tuning больших моделей:

from peft import get_peft_model, LoraConfig, TaskType

config = LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)
model = get_peft_model(model, config)

7. Комбинированные подходы

На практике эффективнее комбинировать методы:

1. Дистилляция → маленькая модель
2. Pruning → удалить неважные веса
3. Квантизация → INT8
4. Компиляция → TensorRT/ONNX

Такой pipeline может дать 50-100x ускорение при приемлемой потере точности.

Выбор метода

• Если нужно быстро: Quantization + Batch norm folding
• Если можно переобучить: Knowledge Distillation
• Если критична точность: Low-rank + Quantization-aware training
• Для production: Комбинировать несколько методов + TensorRT/ONNX

Pro tip: Всегда профилируй модель перед и после оптимизации, проверяй на целевом железе (GPU, Edge device и т.д.). Теоретический спидап может не совпадать с практическим.