При разработке современных приложений искусственного интеллекта эффективный и совместимый запуск обученных моделей машинного обучения в производственных средах является критически важной задачей. Различные фреймворки глубокого обучения (PyTorch, TensorFlow и т.д.) имеют свои собственные стандарты для форматов моделей, что может усложнить кроссплатформенное развертывание и оптимизацию производительности. Именно здесь вступают в игру ONNX (Open Neural Network Exchange) и ONNX Runtime.

Что такое ONNX и почему это важно?

ONNX — это открытый стандарт формата для моделей искусственного интеллекта. Этот стандарт позволяет разработчикам легко переносить модели из одной среды в другую и запускать их в различных аппаратных/программных средах. Конвертируя модель в формат ONNX, мы делаем ее независимой от среды, в которой она была обучена. Это дает значительные преимущества, особенно в следующих сценариях:

• Объединение моделей, разработанных в разных фреймворках ML, в единый формат.
• Портирование моделей на различные цели развертывания, такие как веб, мобильные устройства, периферийные устройства и облако.
• Использование специализированных ускорителей (GPU, TPU, NPU и т.д.) для оптимизации производительности.

ONNX Runtime: высокопроизводительный механизм вывода

ONNX Runtime — это кроссплатформенный, высокопроизводительный механизм вывода, который запускает модели в формате ONNX. Он может интегрироваться с популярными фреймворками, такими как PyTorch и TensorFlow, и предлагает широкий спектр аппаратного ускорения, от центрального процессора (CPU) до графического процессора (GPU) (CUDA, TensorRT) и даже специализированного оборудования (Intel OpenVINO, AMD ROCm). Основные преимущества ONNX Runtime включают:

• Высокая производительность: Включает встроенные оптимизации графов, которые сокращают время вывода и оптимизируют использование ресурсов.
• Кроссплатформенная поддержка: Работает на многих операционных системах, таких как Windows, Linux, macOS, Android, iOS.
• Гибкость: Поддерживает различные аппаратные ускорители (Execution Providers).

Преобразование модели PyTorch в формат ONNX с помощью Python

Преобразование модели PyTorch в ONNX довольно просто. Сначала давайте создадим и обучим небольшую модель PyTorch (или используем предварительно обученную).

import torch
import torch.nn as nn
import torch.optim as optim

# 1. Определение простой модели PyTorch
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2) # 10 входов, 2 выхода
    
    def forward(self, x):
        return self.fc(x)

# Создание модели и загрузка весов (или обучение)
model = SimpleModel()
# Предположим, что вы обучили свою модель или загрузили предварительно обученные веса здесь
# Например, давайте создадим случайные веса
dummy_input = torch.randn(1, 10) # Один образец, 10 признаков

# Перевод модели в режим оценки
model.eval()

# 2. Преобразование модели PyTorch в ONNX
onnx_path = "simple_pytorch_model.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,
    opset_version=11, # Выберите поддерживаемую версию opset
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} # Включить динамический размер пакета
)

print(f"Модель PyTorch успешно преобразована в формат ONNX: {onnx_path}")

Преобразование модели TensorFlow/Keras в формат ONNX с помощью Python

Для преобразования моделей TensorFlow или Keras в ONNX мы можем использовать библиотеку tf2onnx. Возможно, вам потребуется сначала установить эту библиотеку: pip install tf2onnx

import tensorflow as tf
import tf2onnx
from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE

# 1. Определение простой модели Keras
model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(5,)), # 5 входов
    tf.keras.layers.Dense(2, activation='softmax') # 2 выхода
])

# Обучение модели или загрузка весов
# Например, давайте скомпилируем модель
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Проверка вывода модели со случайными данными (необязательно)
dummy_input = tf.random.normal((1, 5))
_ = model(dummy_input) # Убедитесь, что модель строит свою структуру

# 2. Преобразование модели Keras в ONNX
onnx_path = "simple_keras_model.onnx"
input_signature = [tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype, name=model.inputs[0].name)]

model_proto, _ = tf2onnx.convert.from_keras(model, input_signature, opset=11, output_path=onnx_path)

print(f"Модель Keras успешно преобразована в формат ONNX: {onnx_path}")

Выполнение вывода с помощью ONNX Runtime

Давайте посмотрим, как запустить преобразованную модель ONNX с помощью ONNX Runtime. Это применимо к моделям, преобразованным как из PyTorch, так и из TensorFlow.

import onnxruntime as ort
import numpy as np

# Загрузите ранее сохраненную модель ONNX
onnx_model_path = "simple_pytorch_model.onnx" # или "simple_keras_model.onnx"

# Создание сессии ONNX Runtime
# Если вы хотите использовать GPU: providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
# Но вам нужно сначала установить ONNX Runtime для CUDA (pip install onnxruntime-gpu)
sess_options = ort.SessionOptions()
session = ort.InferenceSession(onnx_model_path, sess_options, providers=['CPUExecutionProvider'])

# Получение имен входов и выходов модели
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Создание случайных входных данных (в ожидаемой форме и типе модели)
# Для примера модели PyTorch (1, 10) float32
input_data = np.random.rand(1, 10).astype(np.float32) 
# Для примера модели Keras (1, 5) float32
# input_data = np.random.rand(1, 5).astype(np.float32) 

# Выполнение вывода
outputs = session.run([output_name], {input_name: input_data})

print("Результаты вывода с ONNX Runtime:")
print(outputs[0])

Советы по производительности и лучшие практики

• Правильная версия Opset: Убедитесь, что версия opset, которую вы используете при преобразовании, поддерживается вашей целевой версией ONNX Runtime.
• Динамические формы ввода: Если вашей модели необходимо поддерживать разные размеры пакетов или переменные размеры ввода, укажите это при преобразовании ONNX с помощью параметра dynamic_axes.
• Аппаратное ускорение (Execution Providers): ONNX Runtime поддерживает различные аппаратные ускорители (Execution Providers). Например, использование CUDAExecutionProvider или TensorRTExecutionProvider для графических процессоров NVIDIA может значительно повысить производительность по сравнению с центральным процессором. Не забудьте установить соответствующие библиотеки (onnxruntime-gpu) и указать список providers при создании сессии.
• Оптимизация графов: ONNX Runtime автоматически оптимизирует графы моделей. Однако вы можете еще больше повысить производительность, применяя дополнительные оптимизации или такие методы, как квантование, для конкретных сценариев.
• Проверка модели: Проверьте преобразованную модель ONNX, убедившись, что она выдает тот же результат, что и модель в исходном фреймворке.

Заключение

ONNX и ONNX Runtime — это мощный дуэт, который упрощает развертывание моделей машинного обучения и повышает их производительность. Используя эти инструменты, разработчики могут избавиться от зависимостей моделей, выполнять высокопроизводительный вывод на различных аппаратных платформах и быстрее и эффективнее переносить приложения ИИ в производство. Изучение этих технологий становится незаменимым навыком для разработчиков программного обеспечения и энтузиастов ИИ при создании современных инфраструктур ИИ.