Конвертация модели в ONNX-формат

Конвертация моделей в формат ONNX

ONNX (Open Neural Network Exchange) — открытый формат представления ML-моделей, не привязанный к конкретному фреймворку. Конвертированную модель запускают ONNX Runtime, TensorRT, OpenVINO, CoreML — на CPU, GPU, NPU, мобильных устройствах.

Конвертация HuggingFace модели через Optimum

# Установка
pip install optimum[onnxruntime]

# Экспорт через CLI
optimum-cli export onnx \
  --model bert-base-uncased \
  --task text-classification \
  --opset 17 \
  --device cuda \
  --fp16 \
  ./bert-onnx/

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

# Загрузка и автоматическая конвертация
model = ORTModelForSequenceClassification.from_pretrained(
    "cardiffnlp/twitter-roberta-base-sentiment",
    export=True,
    provider="CUDAExecutionProvider"
)
tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")

# Использование идентично обычной HF модели
inputs = tokenizer("Great product!", return_tensors="pt")
outputs = model(**inputs)

Конвертация PyTorch модели напрямую

import torch

class TextClassifier(torch.nn.Module):
    def __init__(self, vocab_size, num_classes):
        super().__init__()
        self.embedding = torch.nn.EmbeddingBag(vocab_size, 128)
        self.fc = torch.nn.Linear(128, num_classes)

    def forward(self, input_ids, offsets):
        x = self.embedding(input_ids, offsets)
        return self.fc(x)

model = TextClassifier(10000, 3)
model.eval()

dummy_input = (
    torch.randint(0, 10000, (32,)),    # input_ids
    torch.tensor([0, 16])               # offsets
)

torch.onnx.export(
    model,
    dummy_input,
    "text_classifier.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,           # свёртка константных выражений
    input_names=["input_ids", "offsets"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "num_tokens"},
    }
)

Оптимизация ONNX графа

from onnxruntime.transformers import optimizer
from onnxruntime.transformers.fusion_options import FusionOptions

# Автоматическая оптимизация для трансформеров
opt_options = FusionOptions("bert")
opt_options.enable_gelu = True
opt_options.enable_layer_norm = True
opt_options.enable_attention = True
opt_options.enable_skip_layer_norm = True

optimized_model = optimizer.optimize_model(
    "bert.onnx",
    model_type="bert",
    num_heads=12,
    hidden_size=768,
    optimization_options=opt_options,
    opt_level=2,            # 0=нет, 1=базовая, 2=расширенная, 99=все
    use_gpu=True,
    only_onnxruntime=False
)
optimized_model.save_model_to_file("bert_optimized.onnx")

Верификация корректности конвертации

import onnx
import onnxruntime as ort
import numpy as np

# Проверка валидности модели
model = onnx.load("bert_optimized.onnx")
onnx.checker.check_model(model)

# Сравнение вывода с оригинальной PyTorch моделью
pt_model.eval()
with torch.no_grad():
    pt_output = pt_model(**inputs).logits.numpy()

ort_session = ort.InferenceSession("bert_optimized.onnx")
ort_output = ort_session.run(None, {k: v.numpy() for k, v in inputs.items()})[0]

# Проверка численного соответствия
np.testing.assert_allclose(pt_output, ort_output, rtol=1e-3, atol=1e-4)
print("✓ ONNX output matches PyTorch output")

Типичные проблемы конвертации

Dynamic control flow: if len(x) > 0: внутри forward — ONNX не поддерживает динамический control flow. Решение: унификация через masking или конвертация через TorchScript.

Custom operators: операторы без ONNX эквивалента. Решение: регистрация custom op или рефакторинг с использованием стандартных операций.

Dynamic shapes: некоторые операции требуют статических размерностей. Решение: правильное указание dynamic_axes или использование фиксированных размеров с padding.

Numerical precision: накопление ошибок в длинных цепочках операций при FP16. Решение: конвертация в FP32, затем квантизация отдельно.