Инференс в Service Worker и офлайн-сценарии

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

Библиотека поддерживает несколько движков выполнения:

  • WebAssembly (WASM) — кроссплатформенное решение, обеспечивающее совместимость с большинством браузеров.
  • WebGL — ускорение на GPU через браузерные API для графической обработки, позволяющее обрабатывать большие тензоры быстрее, чем на CPU.
  • WebGPU (экспериментально) — современный API для работы с GPU, обеспечивающий более эффективное использование ресурсов видеокарты.

Инициализация ONNX Runtime Web

Для работы с ORT Web необходимо создать объект сессии, который инкапсулирует модель и настройки выполнения. Пример базовой инициализации:

import * as ort from 'onnxruntime-web';

// Создание сессии с указанием бэкенда
const session = await ort.InferenceSession.create('model.onnx', {
  executionProviders: ['wasm'], // 'webgl' или 'webgpu' при необходимости ускорения
});

Ключевое значение здесь имеет executionProviders, поскольку выбор движка напрямую влияет на производительность и совместимость. WASM обеспечивает максимальную кроссплатформенность, WebGL — ускорение на GPU, а WebGPU предлагает новейшие оптимизации для больших моделей.

Подготовка входных данных

Модели ONNX принимают входные данные в виде тензоров ort.Tensor. Конструктор тензора требует указания типа данных, размерности и массива значений:

const inputTensor = new ort.Tensor('float32', new Float32Array([1, 2, 3, 4]), [2, 2]);

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

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

Выполнение инференса

Сессия предоставляет метод run, который принимает объект с ключами, соответствующими именам входов модели:

const feeds = { input: inputTensor };
const results = await session.run(feeds);

Возвращаемый объект results содержит все выходы модели. Доступ к отдельным выходам осуществляется через их имена:

const outputTensor = results['output'];
console.log(outputTensor.data);

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

Использование Service Worker

Для офлайн-приложений и сценариев с предзагрузкой моделей рекомендуется выполнять инференс в Service Worker. Это позволяет:

  • Осуществлять инференс в фоновом потоке, не блокируя интерфейс.
  • Загружать и кэшировать модели через Cache API.
  • Работать в офлайн-режиме без постоянного соединения с сервером.

Пример интеграции с Service Worker:

self.addEventListener('install', async (event) => {
  const cache = await caches.open('model-cache');
  await cache.addAll(['/model.onnx']);
});

self.addEventListener('fetch', async (event) => {
  if (event.request.url.endsWith('model.onnx')) {
    const cache = await caches.open('model-cache');
    const cachedResponse = await cache.match(event.request);
    if (cachedResponse) {
      event.respondWith(cachedResponse);
    }
  }
});

Загрузка модели из кэша Service Worker:

const cache = await caches.open('model-cache');
const response = await cache.match('/model.onnx');
const modelArrayBuffer = await response.arrayBuffer();
const session = await ort.InferenceSession.create(modelArrayBuffer, {
  executionProviders: ['wasm']
});

Оптимизация инференса в браузере

  1. Выбор подходящего провайдера: для мобильных устройств лучше использовать WASM, для десктопов с поддержкой WebGL — WebGL.
  2. Батчинг данных: объединение нескольких входов в один тензор позволяет сократить накладные расходы на вызовы run.
  3. Асинхронная обработка: использование await и Web Worker/Service Worker исключает блокировку UI.
  4. Кэширование моделей: Service Worker и IndexedDB позволяют хранить и переиспользовать модели без повторной загрузки.
  5. Профилирование и warm-up: для крупных моделей рекомендуется выполнять несколько прогонов на пустых данных, чтобы JIT-движок оптимизировал вычисления.

Особенности офлайн-сценариев

  • Полная автономность: весь процесс инференса выполняется на клиенте, данные не отправляются на сервер, что критично для приложений с высокой приватностью.
  • Минимизация загрузки сети: модели хранятся локально, что уменьшает зависимость от сетевого соединения.
  • Контроль версии модели: обновления модели можно осуществлять через Service Worker, сохраняя предыдущие версии для отката.

Использование ONNX Runtime Web вместе с Service Worker открывает возможности для сложных клиентских приложений, где требуется быстрый и безопасный инференс нейросетей, работа в офлайн-режиме и высокая отзывчивость интерфейса.