Кастомные операторы на JavaScript

Для работы с ONNX Runtime Web (ORT Web) используется пакет onnxruntime-web, который доступен через npm. Основной способ подключения выглядит следующим образом:

import * as ort from 'onnxruntime-web';

Возможна также загрузка через CDN:

После подключения создаётся сессия для выполнения модели:

const session = await ort.InferenceSession.create('model.onnx');

Архитектура кастомных операторов

Кастомные операторы (custom operators) позволяют расширять стандартный набор операций ONNX и использовать свои функции внутри модели. Основные компоненты кастомного оператора:

  • Имя оператора — уникальное имя, по которому ORT будет идентифицировать функцию.
  • Версия схемы — версия ONNX, для которой создаётся оператор.
  • Функция вычисления — реализация логики оператора на JavaScript.
  • Сигнатура входов и выходов — определение типов и форм тензоров.

Пример структуры объекта кастомного оператора:

const customOp = {
  name: 'MyCustomOp',
  domain: 'com.example',
  sinceVersion: 1,
  inputs: ['X'],
  outputs: ['Y'],
  typeConstraints: {
    'T': ['tensor(float32)', 'tensor(int32)']
  },
  implementation: (inputs) => {
    const X = inputs.X.data;
    const Y = new Float32Array(X.length);
    for (let i = 0; i < X.length; i++) {
      Y[i] = X[i] * 2; // Простая операция: умножение на 2
    }
    return { Y: new ort.Tensor('float32', Y, inputs.X.dims) };
  }
};

Регистрация кастомного оператора

Для того чтобы ORT распознавал кастомный оператор, его необходимо зарегистрировать перед созданием сессии:

ort.registerCustomOp(customOp);

После регистрации можно загружать модель, содержащую этот оператор, и выполнять инференс.

Типы входов и выходов

Кастомный оператор принимает и возвращает объекты ort.Tensor. Важные моменты:

  • Тип данных: float32, int32, bool и другие стандартные ONNX-типы.
  • Форма тензора (dims): массив, определяющий размерность. Например, [2,3] для матрицы 2x3.
  • Доступ к данным: через свойство .data, которое содержит TypedArray.

Пример использования внутри оператора:

const inputTensor = inputs.X; // ort.Tensor
const data = inputTensor.data; // Float32Array или Int32Array

Оптимизация вычислений

Для больших моделей и сложных операторов рекомендуется:

  1. Использовать TypedArray вместо стандартных массивов для быстрого доступа к числам.
  2. Минимизировать создание новых объектов внутри цикла.
  3. Использовать векторные операции через WebAssembly или SIMD, если требуется высокая производительность.

Обработка нескольких входов и выходов

Кастомный оператор может принимать несколько входных тензоров и возвращать несколько выходных. Синтаксис:

const multiInputOp = {
  name: 'AddMultiply',
  domain: 'com.example',
  sinceVersion: 1,
  inputs: ['A', 'B', 'C'],
  outputs: ['D', 'E'],
  implementation: ({A, B, C}) => {
    const D = new Float32Array(A.data.length);
    const E = new Float32Array(A.data.length);
    for (let i = 0; i < A.data.length; i++) {
      D[i] = A.data[i] + B.data[i];
      E[i] = D[i] * C.data[i];
    }
    return { D: new ort.Tensor('float32', D, A.dims), E: new ort.Tensor('float32', E, A.dims) };
  }
};
ort.registerCustomOp(multiInputOp);

Отладка и тестирование

Для тестирования кастомного оператора можно создать небольшие тензоры и проверить работу функции без полной модели:

const X = new ort.Tensor('float32', Float32Array.from([1, 2, 3]), [3]);
const Y = customOp.implementation({ X });
console.log(Y.Y.data); // Ожидаемый результат: [2, 4, 6]

Использование console.log или библиотек для визуализации массива помогает быстро выявлять ошибки в логике оператора.

Интеграция с асинхронными вычислениями

Если реализация оператора использует асинхронные операции, например, загрузку данных или вызовы WebGPU, функция может возвращать Promise:

const asyncOp = {
  name: 'AsyncOp',
  domain: 'com.example',
  sinceVersion: 1,
  inputs: ['X'],
  outputs: ['Y'],
  implementation: async ({X}) => {
    await fetchSomeData();
    const Y = X.data.map(v => v + 1);
    return { Y: new ort.Tensor('float32', Float32Array.from(Y), X.dims) };
  }
};

ONNX Runtime Web корректно обрабатывает асинхронные кастомные операторы при инференсе.

Совместимость с WebGPU и WebAssembly

ORT Web позволяет запускать вычисления на WebAssembly (по умолчанию) или WebGPU. Кастомные операторы могут быть совместимы с любым бэкендом, но при использовании WebGPU следует учитывать, что операции должны быть выполнены на CPU и возвращены в виде тензора, поскольку прямой доступ к GPU-памяти ограничен.

Оглавление