Оптимизация перед деплоем: onnxoptimizer и onnxsim

ONNX Runtime Web (ORT Web) позволяет запускать модели машинного обучения напрямую в браузере с использованием JavaScript. Для эффективного использования моделей в веб-приложениях критически важно минимизировать размер модели и ускорить время её инициализации. Два ключевых инструмента для этого — onnxoptimizer и onnx-simplifier (onnxsim).

ONNX Optimizer: снижение сложности модели

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

Основные категории оптимизаций:

  • Удаление лишних узлов: устраняет операции, которые не влияют на вычисление (например, Identity или Dropout на этапе инференса).
  • Слияние последовательных операций: объединяет несколько последовательных узлов в один (например, Conv + BatchNorm → Conv).
  • Константное свёртывание (Constant Folding): вычисляет значения выражений, которые полностью определены на этапе компиляции, вместо того чтобы считать их во время инференса.
  • Реализация упрощённых математических трансформаций: заменяет сложные операции на эквивалентные, более эффективные по вычислительной стоимости.

Пример применения в Node.js:

const fs = require('fs');
const onnx = require('onnxjs-node');
const optimizer = require('onnxoptimizer');

const modelBuffer = fs.readFileSync('model.onnx');
const optimizedModel = optimizer.optimize(modelBuffer, [
  'eliminate_identity',
  'fuse_consecutive_transposes',
  'fuse_bn_into_conv'
]);
fs.writeFileSync('model_optimized.onnx', optimizedModel);

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

ONNX Simplifier: упрощение графа модели

onnx-simplifier (или onnxsim) решает задачу упрощения графа после базовой оптимизации, удаляя избыточные вычислительные зависимости и делая модель более предсказуемой для ORT Web.

Преимущества упрощения:

  • Удаление условных и избыточных блоков: упрощает управление потоками данных внутри модели.
  • Выравнивание размеров тензоров: гарантирует, что все промежуточные операции имеют статические и совместимые размеры, что критично для WebAssembly и WebGPU бэкендов ORT Web.
  • Упрощение сложных схем вычислений: объединяет конкатенации, reshape и slice-операции, чтобы минимизировать накладные расходы на браузерные вычисления.

Пример использования в Python (так как onnxsim чаще применяется через Python-интерфейс):

import onnx
from onnxsim import simplify

model = onnx.load('model_optimized.onnx')
model_simplified, check = simplify(model)
if check:
    onnx.save(model_simplified, 'model_simplified.onnx')

Проверка check гарантирует, что упрощённая модель сохраняет точность оригинальной.

Стратегия подготовки модели для ORT Web

  1. Оптимизация веса и структуры

    • Использовать onnxoptimizer для удаления лишних операций и константного свёртывания.
    • Сливать последовательные слои и батч-нормализации.

  2. Стабилизация и упрощение графа

    • Применить onnx-simplifier для унификации размеров тензоров и упрощения ветвлений.
    • Проверить согласованность типов данных и размеров на всех входах и выходах.

  3. Сжатие модели для веб-исполнения

    • Использовать float16 или int8 квантизацию после оптимизации (поддерживается в ONNX).
    • Минимизировать размер файла модели для ускорения загрузки в браузере.

  4. Тестирование на целевом бэкенде

    • ORT Web поддерживает разные бэкенды: wasm, webgl, webgpu.
    • Тестировать упрощённую модель на всех целевых бэкендах, чтобы выявить узкие места производительности.

Практические рекомендации

  • Последовательность оптимизации важна: сначала onnxoptimizer, затем onnx-simplifier.
  • Сохранять промежуточные версии: для анализа влияния каждой оптимизации на производительность.
  • Использовать проверки: встроенные средства ORT Web позволяют сравнить предсказания исходной и оптимизированной моделей.
  • Минимизировать ветвления и условные операции: браузерные движки хуже обрабатывают динамические графы, что может привести к падению производительности.

Интеграция в процесс сборки веб-приложения

Оптимизация моделей должна быть частью CI/CD процесса. Скрипт автоматической оптимизации:

  1. Загружает исходный ONNX-модель.
  2. Применяет onnxoptimizer с заранее выбранными правилами.
  3. Упрощает модель через onnxsim.
  4. Применяет квантизацию и сохраняет финальную версию в папку с ассетами веб-приложения.

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

Оглавление