Vite: конфигурация, заголовки COOP/COEP через плагины

ONNX Runtime Web (ORT Web) предоставляет возможность выполнять инференс моделей ONNX прямо в браузере с использованием JavaScript или TypeScript. Для начала работы требуется установка пакета через npm или yarn:

npm install onnxruntime-web

После установки библиотека импортируется стандартным способом:

import * as ort from 'onnxruntime-web';

Создание сессии для модели происходит с использованием InferenceSession.create, где указывается путь к файлу модели и опции сессии:

const session = await ort.InferenceSession.create('model.onnx', {
  executionProviders: ['wasm'], // Возможные значения: 'wasm', 'webgl', 'cpu'
});

Ключевые моменты:

  • Execution Providers определяют, как будет выполняться инференс: в CPU через WASM или с ускорением WebGL.
  • Для большинства моделей рекомендуется сначала тестировать wasm, затем при необходимости переходить на webgl для ускорения вычислений на GPU.

Подготовка данных и инференс

Данные для модели в ORT Web передаются в формате ort.Tensor. Создание тензора требует указания типа данных, формы и массива значений:

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

Запуск инференса осуществляется через метод session.run, который принимает объект с именами входов:

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

Особенности работы:

  • Результаты возвращаются в виде объекта { output_name: ort.Tensor }.
  • Форма и тип выходного тензора соответствуют описанию модели.

Интеграция в Vite

Vite требует правильной конфигурации для работы с WASM и WebGL модулями, которые использует ONNX Runtime Web. В vite.config.js следует включить поддержку ассетов и настроить плагины для правильной обработки файлов .wasm.

import { defineConfig } from 'vite';

export default defineConfig({
  build: {
    target: 'esnext',
    assetsInclude: ['**/*.wasm']
  }
});

Важно:

  • Файлы .wasm должны попадать в сборку как ассеты, иначе при загрузке сессии возникнет ошибка.
  • Использование target: 'esnext' позволяет избежать проблем с асинхронной загрузкой WASM в старых сборках.

Настройка заголовков COOP и COEP через плагины

Для корректной работы WebAssembly и модулей WebGL необходимо включить заголовки Cross-Origin-Opener-Policy (COOP) и Cross-Origin-Embedder-Policy (COEP). Это обеспечивает безопасное взаимодействие с памятью и доступ к SharedArrayBuffer.

Пример конфигурации плагина для Vite:

import { defineConfig } from 'vite';
import vitePluginHeaders from 'vite-plugin-headers';

export default defineConfig({
  plugins: [
    vitePluginHeaders({
      '/**': {
        'Cross-Origin-Opener-Policy': 'same-origin',
        'Cross-Origin-Embedder-Policy': 'require-corp'
      }
    })
  ]
});

Особенности:

  • Заголовки должны применяться ко всем маршрутам ('/**'), чтобы SharedArrayBuffer был доступен в любом скрипте.
  • COOP = same-origin изолирует контекст окна, предотвращая утечки данных между сайтами.
  • COEP = require-corp гарантирует, что встраиваемые ресурсы приходят с заголовком Cross-Origin-Resource-Policy: same-origin или Cross-Origin-Resource-Policy: cross-origin.

Оптимизация загрузки моделей

Модели ONNX могут быть крупными, поэтому для уменьшения времени загрузки и снижения потребления памяти рекомендуется:

  1. Асинхронная загрузка:
const sessionPromise = ort.InferenceSession.create('model.onnx', { executionProviders: ['wasm'] });
// инференс выполняется позже
  1. Кэширование сессий:

    • Один раз созданная сессия может использоваться многократно для разных данных.
  2. Lazy загрузка ресурсов:

    • WASM-файлы загружаются только при необходимости инференса.

Использование WebGL для ускорения инференса

Для моделей с большим количеством операций и слоев рекомендуется WebGL:

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

Преимущества:

  • Графические ускорители обрабатывают матричные вычисления быстрее CPU.
  • Снижается нагрузка на основной поток JavaScript, улучшая отзывчивость интерфейса.

Ограничения:

  • Поддерживаются не все операции ONNX.
  • Требуется включение COOP/COEP для работы SharedArrayBuffer в браузерах с усиленной политикой безопасности.

Работа с несколькими входами и выходами

Если модель имеет несколько входов:

const input1 = new ort.Tensor('float32', new Float32Array([1,2,3]), [3]);
const input2 = new ort.Tensor('float32', new Float32Array([4,5,6]), [3]);

const results = await session.run({
  'input1_name': input1,
  'input2_name': input2
});

Результаты можно использовать сразу по именам выходов:

const output1 = results['output1_name'].data;
const output2 = results['output2_name'].data;

Диагностика ошибок и отладка

Типичные ошибки:

  • Ошибка загрузки WASM: возникает при отсутствии файла .wasm в сборке Vite.
  • SharedArrayBuffer недоступен: необходимо включить COOP/COEP заголовки.
  • Несоответствие форматов данных: входной тензор должен строго соответствовать форме и типу модели.

Рекомендуется выводить ошибки в консоль и проверять структуру входных и выходных тензоров для отладки инференса.