ONNX и ONNX Runtime: как работает вывод моделей и оптимизация

ONNX (Open Neural Network Exchange) — это общий файловый формат обученной модели: он включает график вычислений и схему с сохранёнными весами. ONNX Runtime (ORT) от Microsoft — это механизм логического вывода, который позволяет запускать ONNX-модели на высокой скорости, причём кроссплатформенно.

Рассмотрим несколько вопросов и ответов о работе с ONNX и ORT.

1. В чём разница между ONNX и ONNX Runtime?
   ONNX — это формат файла (данные), а ONNX Runtime — механизм, который этот файл запускает. Например, в работе с rembg (инструментом для удаления фона) ONNX выступает в роли чертежа (blueprint), а ORT — в роли движка.

2. Почему переход на PyTorch может быть проще и быстрее, если общий формат вроде бы работает медленно?
   Обычное выполнение в PyTorch (eager) динамически обрабатывается с помощью Python: за одну операцию выполняется ряд задач, включая автоматическую дифференциацию (autograd) для обучения. Это обеспечивает гибкость, но влечёт высокие накладные расходы. В свою очередь, ONNX представляет всю модель в виде фиксированного графика, поэтому ORT можно предварительно оптимизировать: объединить арифметические операции (например, Conv + ReLU в одну), выполнить постоянную свёртку, удалить ненужные узлы, оптимизировать размещение в памяти. При этом в ORT отсутствует автоматическая дифференциация и цикл Python.

Иными словами, общий формат сам по себе не замедляет выполнение — источник скорости в том, что модель можно предварительно оптимизировать благодаря её статичности. Медленность PyTorch во время выполнения связана с его гибкостью, а ORT эту проблему устраняет. При этом PyTorch тоже поддерживает работу с ORT — не всегда ORT выигрывает, иногда быстрее оказывается компиляция и т. д. Для процессора дистрибутив с ORT особенно экономичен.

1. Провайдер выполнения работает автоматически или вручную?
   Это гибрид: сначала вручную задаётся список приоритетов, а затем система автоматически назначает каждый узел. Если вы задаёте кандидатуру и приоритет для использования, ORT присваивает каждой операции соответствующий EP с наивысшим приоритетом, а неподдерживаемую операцию автоматически возвращает в CPU. Версия ORT для CPU использует только CPU EP, а версия for GPU может задействовать CPU, GPU и CUDA.

2. Насколько сильно снижается точность квантования? В каких случаях его стоит применять?
   Стандартное квантование — это переход от fp32 к int8: размер данных сокращается примерно в 4 раза, а скорость работы процессора увеличивается в 2–4 раза. При этом снижение точности во многих CNN составляет от 1 % до нескольких % — оно зависит от модели и задачи.

Есть два подхода к квантованию:
- динамическое (только для веса int8) — простой вариант, ориентирован на NLP и трансформеры;
- статическое (требует калибровки) — обеспечивает высокую точность при работе с CNN.

Квантование стоит использовать, когда нужно оптимально распределить ресурсы процессора или работать с граничными условиями — если допустимо небольшое снижение точности. Избегайте его в приложениях, где требуется высокая точность: например, тонкие участки (вроде границ волосков при удалении фона) могут быть сглажены.

1. В чём подробные различия и как обеспечить совместимость с PyTorch (например, как снимать слои)? Подход к детализации (сниманию слоёв) в PyTorch и ONNX существенно различается.

В PyTorch используется иерархия модулей (например, Conv2d или вложенных самодельных классов): график определяется во время выполнения (define-by-run).

В ONNX иерархия расширяется до плоского графа операторов — набора узлов (например, Conv, ReLU, Add). Здесь нет концепции «слой/класс» высокого уровня — только узлы и рёбра.

С помощью torch.onnx.export можно отследить и обработать код PyTorch, преобразовав его в статический график. Но для совместимости нужно учитывать несколько моментов:
- версия opset должна соответствовать;
- если в модели есть операции, не поддерживаемые ONNX (например, самодельные операции, динамический поток управления, специальная динамическая форма), экспорт может завершиться ошибкой — такие операции придётся избегать или модифицировать.

ORT специализируется в первую очередь на логическом выводе (хотя существует и обучающий ORT — это отдельная история).

1. Ускорится ли версия CPU за счёт NEON/SIMD? Да, это происходит автоматически. При выполнении ORT на процессоре используется библиотека MLAS — она обнаруживает и задействует инструкции CPU SIMD во время работы.
2. для x86: AVX2 / AVX-512;
3. для ARM (например, Apple Silicon в смартфонах): NEON.

Во время квантования int8 тоже задействуются специальные инструкции для ускорения операций с точечным продуктом (например, VNNI).

Кроме того, можно настроить количество потоков (OMPNUMTHREADS, intra/interop-op) и уровень оптимизации графика (ORTENABLEALL). Например, rembg так быстро работает на Mac именно благодаря использованию NEON.

Таким образом, ONNX — это общий формат модели, а ONNX Runtime — механизм логического вывода. Скорость работы достигается за счёт предварительной оптимизации статического графика. ONNX можно запускать на CPU, GPU и разных ОС, а SIMD-инструкции (AVX, NEON) автоматически ускоряют работу на процессоре. Именно поэтому, к примеру, rembg смог работать на процессоре без PyTorch.

Эта статья подготовлена на основе диалога с генеративным ИИ. Также с помощью ИИ была создана диаграмма. Ознакомьтесь с последними спецификациями на официальном сайте ONNX Runtime.