Ускоряем инференс в Python с ONNX

Привет! Если у вас когда‑либо был опыт деплоя нейросетки, вы знаете, что обучение — это полдела, а вот добиться шустрого инференса — целое искусство. Часто обученная в PyTorch модель дает замечательные метрики, но стоит попытаться запустить её в приложении начинаются всякие проблемки.

Одно из решений, которое часто выручает — ONNX и ONNX Runtime. Если эти буквы для вас пока ничего не значат — не беда, сейчас разберёмся что к чему. А если вы уже слышали про ONNX, то, возможно, задавались вопросом: «А реально ли ускорить инференс, заморочившись с этой технологией?» Еще как!

Что такое ONNX и зачем он нужен

Начнём с базы. ONNX (Open Neural Network Exchange) — это открытый формат для представления нейронных сетей. Такой вот способ сохранить модель не привязанной к конкретному фреймворку. Мы можем обучить модель хоть в PyTorch, хоть в TensorFlow, экспортировать её в ONNX — и затем запускать где угодно: в Python, в Java, в JavaScript, на мобильном, на C++… Кросс‑платформенность и интероперабельность — главные фичи ONNX.

Однако сегодня нас интересует не столько переносимость, сколько производительность. Тут поможет ONNX Runtime — высокопроизводительный движок для выполнения моделей в формате ONNX. ONNX Runtime умеет всякие умные штуки: оптимизацию графа вычислений, слияние операций, использование аппаратных ускорителей.

Если правильно готовить, модель в ONNX Runtime может работать значительно быстрее, чем в исходном фреймворке. Особенно это заметно на CPU, где динамические вычисления PyTorch иногда проигрывают оптимизированному под конкретную модель графу.

Прикинем план действий:

1. Берём обученную модель в PyTorch (научим или загрузим — не суть).

2. Экспортируем её в файл .onnx.

3. Загружаем этот файл с помощью ONNX Runtime в Python.

4. Гоним инференс и измеряем скорость, сравниваем с оригинальным PyTorch.

5. Profit!.

Нанчем.

Экспорт модели из PyTorch в ONNX

Допустим есть нейросеть, обученная с помощью PyTorch. Для примера возьмём что‑нибудь распространённое, например ResNet-18 — классическую сверточную сеть для классификации изображений. У вас это может быть своя модель — порядок действий не меняется.

Первым делом, нужно сохранить модель в формате ONNX. В PyTorch это делается через функцию torch.onnx.export. Она вынимает структуру модели и текущие веса и складывает их в единый файл. Код будет примерно таким:

import torch
from torchvision import models

# Загружаем предобученную модель ResNet-18 (для примера)
model = models.resnet18(pretrained=True)
model.eval()  # перевести в режим инференса

# Подготовим тестовый вход - тензор нужной формы
dummy_input = torch.randn(1, 3, 224, 224)  # батч из 1 изображения 3x224x224

# Экспорт модели в ONNX
torch.onnx.export(
    model,               # модель для конвертации
    dummy_input,         # пример входных данных (для отслеживания размеров)
    "resnet18.onnx",     # имя выходного файла
    opset_version=13,    # версия ONNX opset, 13 – довольно универсально
    input_names=["input"],   # имена входных тензоров (необязательно)
    output_names=["output"], # имена выходных (необязательно)
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}} 
    # dynamic_axes позволяет задать изменяемый размер батча
)
print("Модель успешно экспортирована в resnet18.onnx")

Указали opset_version=13 — это версия спецификации ONNX, поддерживаемая большинством фреймворков (актуальные версии выше, но 13 — надёжный выбор на 2023–2024 годы). Также dynamic_axes задаёт, что нулевой размер (батч) у входа и выхода может меняться — чтобы мы не были зафиксированы на batch_size=1.

После выполнения этого кода на диске появится файл resnet18.onnx. Это уже не PyTorch‑модель, а граф вычислений с весами, который можно тягать куда угодно. Весит он, кстати, примерно столько же, сколько PyTorch‑снапшот, может чуть больше, ведь там хранится та же информация плюс немного служебной.

Инференс с помощью ONNX Runtime

Теперь запустим модель через ONNX Runtime. Предварительно нужно установить пакет onnxruntime (или onnxruntime-gpu, если хотите на GPU, об этом позже). Устанавливается просто: pip install onnxruntime. Есть также версия onnxruntime-silicon для Mac (Apple Silicon) и прочие, но базовая onnxruntime автоматом тянет нужное под вашу платформу.

Дальше код:

import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms

# Загружаем нашу ONNX модель
session = ort.InferenceSession("resnet18.onnx", providers=['CPUExecutionProvider'])

# Подготовим функцию загрузки и препроцессинга изображения (как у PyTorch модели)
def load_image(path):
    img = Image.open(path).convert('RGB')
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    tensor = transform(img).unsqueeze(0)  # добавляем batch dimension
    return tensor

# Пример: грузим изображение
input_image = load_image("cat.jpg")  # допустим, фотка котика
input_array = input_image.numpy()

# Получаем имя входного тензора для сессии (как мы указали "input")
input_name = session.get_inputs()[0].name

# Запускаем инференс
output = session.run(None, {input_name: input_array})

# Результат – список из выходных массивов, у нас один выход "output"
output_tensor = output[0]
predictions = output_tensor.flatten()
pred_class = np.argmax(predictions)
print(f"Предсказанный класс: {pred_class}")

Создаём InferenceSession и загружаем в неё модель. Параметром providers указываем CPU. Если установить onnxruntime-gpu, по умолчанию будет пробоваться CUDA. Но давайте пока про CPU.

Мы подготовили картинку так же, как делали бы для PyTorch: привели к размеру 224×224, нормализовали, превратили в тензор. Единственное отличие, для ONNX Runtime надо дать numpy‑массив, а не PyTorch‑тензор, поэтому вызываем .numpy().

Узнаём имя входа: в ONNX модели входы/выходы по имени, можно было самому знать («input»), но на всякий случай я показал, как выдернуть через session.get_inputs().

Вызываем session.run(None, {input_name: input_array}). Первый аргумент None означает «дай мне все выходы». Можно перечислить имена нужных выходов, но у нас один выход, так что не принципиально.

Получаем результат — обычный numpy массив с предсказаниями. Дальше делаем argmax — и получаем индекс класса, который модель считает самым вероятным.

Для проверки можно сравить с выводом оригинальной модели в PyTorch, результаты должны совпадать (до погрешности округления). Если не совпадают, возможно, вы забыли какой‑то препроцессинг или пост‑обработку.

Замеряем скорость

Теперь ради чего всё затевалось. Сравним производительность. Обычно это делают прогоняя несколько сотен или тысяч прогонов и усредняя время. Можно воспользоваться модулем time или timeit. Приведу пример простого замера для нашей модели:

import time

# Функция замера времени для PyTorch
def benchmark_pytorch(model, input_tensor, steps=100):
    model.eval()
    timings = []
    with torch.no_grad():
        for i in range(steps):
            start = time.time()
            _ = model(input_tensor)
            end = time.time()
            timings.append(end - start)
    return sum(timings) / len(timings) * 1000  # мс на шаг

# Функция замера для ONNX Runtime
def benchmark_onnx(session, input_name, input_array, steps=100):
    timings = []
    for i in range(steps):
        start = time.time()
        session.run(None, {input_name: input_array})
        end = time.time()
        timings.append(end - start)
    return sum(timings) / len(timings) * 1000  # мс на шаг

# Прогреем модель PyTorch (чтобы JIT размялся, если применимо) и ONNX
_ = model(input_image)  # прогон для PyTorch
session.run(None, {input_name: input_array})  # прогон для ONNX

# Теперь меряем
pt_time = benchmark_pytorch(model, input_image)
onnx_time = benchmark_onnx(session, input_name, input_array)

print(f"Среднее время инференса PyTorch: {pt_time:.2f} ms")
print(f"Среднее время инференса ONNX: {onnx_time:.2f} ms")
print(f"Ускорение: {pt_time/onnx_time:.2f}x")

Вывод:

Среднее время инференса PyTorch: 4.50 ms  
Среднее время инференса ONNX: 1.80 ms  
Ускорение: 2.50x

Результат для ResNet-18 на одном CPU‑потоке.

Откуда берётся ускорение? ONNX Runtime делает несколько вещей лучше, чем голый PyTorch на CPU:

• Он использует оптимизации уровня графа. Например, может слить последовательность из нескольких операций в одну, убрать лишние слои, предвычислить константы (кстати, PyTorch при экспорте с do_constant_folding=True уже некоторые вещи сглаживает).

• Он задействует высокопроизводительные библиотеки: MKL‑DNN (oneDNN) для x86, или собственные оптимизации Microsoft. PyTorch тоже умеет в MKL, но ONNXRuntime порой эффективнее планирует вычисления.

• Если у вас несколько ядер CPU, ONNXRuntime может выполнять узлы графа параллельно (в примере выше мы явно не делали параллелизм, но опция есть). PyTorch в режиме eval тоже параллелит потоки для некоторых операций, но, опять же, раз на раз.

• В ONNX Runtime легко включить квантизацию модели — буквально пару строчек с использованием ONNX Runtime Tools, и ваша модель перейдёт, например, с FP32 на INT8, что часто даёт ускорение на CPU без большой потери качества. В PyTorch тоже есть квантизация, но её нужно настроить перед экспортом.

Но вообще, не всегда ONNX Runtime гарантирует выигрыш. В большинстве сценариев для CPU — да, особенно на более старых версиях PyTorch, когда TorchScript был сыроват. Сейчас PyTorch 2.x с компиляцией (TorchInductor) тоже может удивить скоростью.

А вот на GPU не всё так однозначно. Если вы просто загрузите модель ONNX и будете гонять её на CUDA через ONNX Runtime, вы можете не получить никакого прироста, а то и замедление.

Чтобы ускорить GPU‑инференс, часто пользуются ONNX совместно с TensorRT — это такая спецоптимизация от NVIDIA, которая умеет из ONNX‑модели сделать сверхоптимизированный бинарный исполнимый план под конкретную видеокарту.

ONNXRuntime может работать с TensorRT как провайдером. Вот там да, можно выжать 2x-4x ускорение на GPU.

Если у вас маленькая модель или и так всё укладывается в требования, может и не стоить городить конвертацию, добавится сложность в пайплайн. Но когда секунды вам важны, почему бы не попробовать? В худшем случае, потратив пару часов, вы убедитесь, что выигрыш невелик. В лучшем получите ускорение x2, x3, а то и x5.

Если ONNX и бенчмарки для вас уже не «магия», а рабочий инструмент, дальше упираетесь в более сложные вопросы: пайплайны, качество, воспроизводимость, ограничения AutoML и выбор методов под задачу. На курсе Machine Learning. Advanced это разбирают на практике — от production-кода и окружения до продвинутых моделей и кейсов уровня Middle+. Чтобы узнать, подойдет ли вам программа курса, пройдите вступительный тест.

Для знакомства с форматом обучения и экспертами приходите на демо-урок «Поиск аномалий во временных рядах: за рамками трех сигм» 18 февраля в 20:00. Участие бесплатное, нужно только зарегистрироваться.

Еще больше бесплатных уроков от преподавателей курсов можно посмотреть в календаре мероприятий.