Введение в ONNX и экспорт моделей из PyTorch

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

Что такое ONNX?

ONNX представляет собой открытый стандарт, который был разработан совместно крупнейшими игроками в области технологий, такими как Microsoft, Amazon и Facebook, для создания открытого экосистема обмена моделями машинного обучения. Это позволяет разработчикам использовать модели, обученные в одном фреймворке (например, PyTorch или TensorFlow), и легко переносить их в другую среду выполнения или на другую платформу без необходимости переписывания кода или переобучения моделей.

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

Установка зависимостей

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

pip install torch onnx onnxruntime

Создание модели

Следующий шаг — подготовка модели, которую вы хотите экспортировать. В примере ниже мы создаем простейшую нейронную сеть для классификации изображений:

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleImageClassifier(nn.Module):
    def __init__(self):
        super(SimpleImageClassifier, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(this.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = SimpleImageClassifier()

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

PyTorch предлагает удобные инструменты для перевода моделей в формат ONNX. Важно обеспечить корректность входных данных модели для успешного экспорта:

dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(model, dummy_input, "model.onnx", export_params=True)

Визуализация и выполнение модели ONNX

После экспорта модели в ONNX формат, возможно визуализировать архитектуру сети, а также осуществить ее запуск:

Визуализация графа

Инструменты вроде Netron или Graphviz позволяют наглядно исследовать структуру экспортированной модели, что может быть полезно для проверки корректности перевода модели в ONNX.

Выполнение модели с ONNX Runtime

ONNX Runtime, разработанный Microsoft, представляет собой производительный движок для выполнения предсказаний с использованием моделей ONNX. Это мощное решение для разнообразных платформ и устройств:

import onnxruntime as ort

# Создание сессии ONNX Runtime
ort_session = ort.InferenceSession("model.onnx")
inputs = ort_session.get_inputs()
outputs = ort_session.get_outputs()

# Подготовка входных данных и выполнение модели
input_data = np.random.randn(1, 1, 28, 28).astype(np.float32)
output = ort_session.run([outputs[0].name], {inputs[0].name: input_data})

Оптимизация моделей ONNX

Оптимизация модели после ее экспорта в ONNX формат важна для улучшения производительности и уменьшения нагрузки на ресурсы. К таким методам оптимизации относятся квантизация и устранение неэффективных операций.

Квантизация моделей

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

from optimum.onnxruntime import ORTConfig, ORTQuantizer

model_ckpt = "your_model_checkpoint"
ort_config = ORTConfig(quantization_approach="dynamic")
quantizer = ORTQuantizer(ort_config)
quantizer.export(model_ckpt, output_path="model-quantized.onnx", feature="sequence-classification")

Экспорт, визуализация, выполнение и оптимизация моделей через формат ONNX открывает большие возможности для разработчиков и компаний, стремящихся максимально эффективно использовать ресурсы и обеспечить быструю работу своих моделей машинного обучения.
Подпишитесь на наш Telegram-канал

Удаление избыточных операций

Помимо квантизации, ещё одним важным аспектом оптимизации моделей в ONNX является удаление избыточных операций. Этот процесс включает в себя анализ графа модели и исключение узлов или операций, которые не влияют на конечный результат. Это может существенно сократить время выполнения модели и уменьшить её размер.

Использование инструментов оптимизации

ONNX предоставляет ряд встроенных возможностей для оптимизации моделей, таких как Model Optimizer, который является частью ONNX Runtime. Этот инструмент автоматически применяет описанные оптимизации, улучшая производительность модели без ущерба для её точности.

from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions

# Настройка опций сессии для использования оптимизаций
options = SessionOptions()
options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_EXTENDED

# Загрузка оптимизированной модели
session = InferenceSession("path_to_model.onnx", options)

Практическое применение ONNX в индустрии

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

Кейс использования ONNX в здравоохранении

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

Применение в автономном вождении

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

Заключение

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

Полезные ссылки

• Официальный сайт ONNX
• ONNX GitHub
• Конвертация моделей в ONNX
• Netron – Визуализатор моделей ONNX

Подпишитесь на наш Telegram-канал