Архитектура ONNX Runtime Web: WebAssembly и WebGL бэкенды

ONNX Runtime Web (ORT Web) представляет собой высокопроизводительную среду выполнения моделей машинного обучения в браузере, разработанную на основе ONNX Runtime. Основная цель ORT Web — обеспечение быстрого и эффективного инференса моделей ONNX непосредственно в клиентской части без необходимости обращения к серверу. Достижение этой цели обеспечивается гибкой архитектурой, включающей поддержку двух ключевых бэкендов: WebAssembly (WASM) и WebGL.

Основные компоненты

1. Ядро ONNX Runtime

Ядро ORT Web выполняет функции загрузки, инициализации и управления моделью. Оно абстрагирует детали работы конкретного бэкенда и предоставляет единый интерфейс для работы с входными и выходными тензорами.

2. Бэкенды исполнения

  • WebAssembly (WASM): основной бэкенд для большинства задач. Использует компиляцию кода ONNX Runtime на уровень WASM, что обеспечивает переносимость и высокую производительность на широком спектре устройств.
  • WebGL: GPU-ускоренный бэкенд, предназначенный для параллельного вычисления тензоров с высокой размерностью. Использует графический API браузера для выполнения операций на видеокарте, что особенно эффективно для свёрточных нейросетей и глубокого обучения.

WebAssembly бэкенд

Архитектура и принципы работы

WebAssembly бэкенд реализует стандартные операции ONNX через интерпретацию или компиляцию WASM-модулей. Основные особенности:

  • Кроссплатформенность: поддержка всех современных браузеров без дополнительных плагинов.
  • Оптимизация памяти: управление heap-областью позволяет избежать частых копирований данных и снижает задержки при инференсе.
  • Синхронное и асинхронное выполнение: позволяет использовать как блокирующие, так и неблокирующие операции, что важно для интеграции в интерактивные веб-приложения.

Принцип работы

  1. Загрузка WASM-модуля ORT через fetch и инициализация с помощью WebAssembly.instantiate.
  2. Создание сессии с моделью ONNX через InferenceSession.create().
  3. Передача входных тензоров в сессию.
  4. Получение результатов инференса в виде тензоров и преобразование их в удобные форматы (TypedArray или объекты JS).

WebGL бэкенд

Особенности и возможности

WebGL бэкенд использует графический процессор для вычислений, что обеспечивает значительное ускорение при работе с матрицами высокой размерности:

  • Параллельные вычисления: каждая операция на тензорах переводится в шейдерные программы, которые выполняются параллельно на GPU.
  • Поддержка больших моделей: эффективная работа с моделями, требующими значительных вычислительных ресурсов, например, ResNet или Transformer.
  • Оптимизация графа вычислений: библиотека строит внутренний граф операций и применяет техники слияния слоёв и буферного переиспользования.

Принцип работы

  1. Инициализация WebGL контекста через HTMLCanvasElement или OffscreenCanvas.
  2. Преобразование входных данных в текстуры GPU.
  3. Выполнение операций через GLSL-шейдеры, каждая из которых соответствует одной или нескольким операциям ONNX.
  4. Сбор результатов из текстур и преобразование их обратно в тензоры JavaScript.

Управление сессиями и оптимизация

Сессии ORT Web являются центральным элементом для инференса. Они предоставляют методы:

  • session.run(feeds): выполнение инференса с передачей словаря входных данных.
  • session.end(): освобождение ресурсов, связанных с бэкендом.
  • session.getInputNames() и session.getOutputNames(): получение метаданных модели.

Оптимизация производительности

  • Использование WebAssembly SIMD позволяет ускорить операции на CPU за счет векторизации.
  • WebGL бэкенд может использовать текстурные буферы большого размера, минимизируя количество операций передачи данных между CPU и GPU.
  • Поддержка multi-threading через Web Workers для выполнения нескольких инференсов параллельно без блокировки основного потока UI.

Совместное использование бэкендов

ORT Web позволяет динамически выбирать бэкенд в зависимости от характеристик устройства и типа модели. Например, для малых моделей предпочтителен WASM для минимизации накладных расходов, а для сложных свёрточных сетей — WebGL для максимального ускорения.

Выбор бэкенда производится при инициализации сессии:

const session = await ort.InferenceSession.create(modelUrl, { executionProviders: ['webgl', 'wasm'] });

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

Особенности работы с тензорами

ORT Web использует собственный тип данных Tensor, совместимый с TypedArray JavaScript. Особенности:

  • Поддержка различных типов: float32, int32, uint8, bool.
  • Возможность прямого создания из существующих TypedArray без копирования.
  • Методы для конвертации между форматами и управления размерами (shape) и типами данных.

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

ORT Web легко интегрируется с популярными фронтенд-фреймворками, обеспечивая:

  • Асинхронную загрузку моделей и данных.
  • Минимизацию блокировок UI за счет Web Workers.
  • Поддержку мобильных устройств и браузеров с ограниченными ресурсами.

Архитектура с чётким разделением ядра, бэкендов и сессий обеспечивает расширяемость и высокую адаптивность ORT Web для любых веб-приложений с задачами машинного обучения.

Оглавление