Использование кастомных операторов

ONNX Runtime Web (ORT Web) предоставляет возможность выполнять модели машинного обучения в браузере или Node.js, используя стандартные и кастомные операторы. Кастомные операторы необходимы для реализации функциональности, которая отсутствует в стандартном наборе ONNX, или для оптимизации специфичных вычислительных задач.

Подключение и регистрация кастомных операторов

Для работы с кастомными операторами в ORT Web используется объект Ort.CustomOp. Каждое определение оператора включает имя, количество входов и выходов, типы данных, а также функцию вычисления.

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

const customOp = {
    name: 'MyCustomOp',
    inputTypes: ['float32', 'float32'],
    outputTypes: ['float32'],
    kernel: (inputs, outputs) => {
        const [a, b] = inputs;
        const result = outputs[0];
        for (let i = 0; i < a.length; i++) {
            result[i] = a[i] + b[i]; // элементное сложение
        }
    }
};

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

  • name — уникальное имя оператора, которое будет использоваться в модели ONNX.
  • inputTypes и outputTypes — массивы типов данных для входов и выходов. Поддерживаются стандартные типы ONNX, такие как float32, int32, bool.
  • kernel — функция, реализующая вычислительную логику оператора. Она получает массивы входов и выходов в виде TypedArray.

Для регистрации кастомного оператора используется метод Ort.env.registerCustomOp:

Ort.env.registerCustomOp(customOp);

После регистрации модель, содержащая оператор MyCustomOp, сможет корректно выполняться в браузере или Node.js без ошибок.

Работа с несколькими кастомными операторами

ORT Web позволяет регистрировать сразу несколько операторов одновременно. Для этого используется массив объектов:

const customOps = [
    {
        name: 'AddScalar',
        inputTypes: ['float32'],
        outputTypes: ['float32'],
        kernel: (inputs, outputs, attributes) => {
            const scalar = attributes.scalar || 0;
            const input = inputs[0];
            const output = outputs[0];
            for (let i = 0; i < input.length; i++) {
                output[i] = input[i] + scalar;
            }
        }
    },
    {
        name: 'Square',
        inputTypes: ['float32'],
        outputTypes: ['float32'],
        kernel: (inputs, outputs) => {
            const input = inputs[0];
            const output = outputs[0];
            for (let i = 0; i < input.length; i++) {
                output[i] = input[i] * input[i];
            }
        }
    }
];

customOps.forEach(op => Ort.env.registerCustomOp(op));

Особенности использования:

  • Атрибуты оператора передаются как объект attributes в функцию kernel.
  • Все вычисления выполняются на стороне JavaScript, что позволяет использовать возможности TypedArray для высокой производительности.

Интеграция кастомных операторов в модель

Чтобы модель ONNX корректно использовала кастомный оператор, он должен быть определён в графе модели с соответствующим именем и набором входов/выходов. В ONNX можно создавать кастомные узлы через NodeProto или использовать утилиты типа onnx.helper.make_node.

Пример создания узла с кастомным оператором через Python и экспорт в ONNX:

import onnx
from onnx import helper, TensorProto

node = helper.make_node(
    'MyCustomOp',
    inputs=['input1', 'input2'],
    outputs=['output1']
)

graph = helper.make_graph(
    nodes=[node],
    name='CustomOpGraph',
    inputs=[
        helper.make_tensor_value_info('input1', TensorProto.FLOAT, [None]),
        helper.make_tensor_value_info('input2', TensorProto.FLOAT, [None])
    ],
    outputs=[
        helper.make_tensor_value_info('output1', TensorProto.FLOAT, [None])
    ]
)

model = helper.make_model(graph)
onnx.save(model, 'custom_op_model.onnx')

После загрузки в ORT Web модель будет использовать зарегистрированный кастомный оператор.

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

  1. Использование TypedArray: все входные и выходные данные передаются как Float32Array, Int32Array и т.д., что минимизирует накладные расходы на преобразование данных.
  2. Векторизация вычислений: внутри kernel рекомендуется использовать циклы по элементам массива или WebAssembly/ SIMD-расширения для ускорения обработки больших данных.
  3. Минимизация аллокаций: выделение массивов заранее снижает нагрузку на сборщик мусора.

Особенности отладки

  • ORT Web позволяет отслеживать ошибки кастомного оператора через стандартные механизмы try/catch.
  • Для сложных вычислений можно временно логировать промежуточные массивы или размеры тензоров.
  • Если модель падает с ошибкой типа Node missing kernel for custom op, это обычно означает, что оператор не зарегистрирован или имя не совпадает с определением в модели.

Поддержка атрибутов

Кастомные операторы могут принимать атрибуты, передаваемые из модели. В функции kernel они доступны как объект attributes. Атрибуты могут быть числами, строками или массивами, что позволяет гибко настраивать поведение оператора без изменения кода JavaScript.

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

const multiplyBy = {
    name: 'MultiplyBy',
    inputTypes: ['float32'],
    outputTypes: ['float32'],
    kernel: (inputs, outputs, attributes) => {
        const factor = attributes.factor ?? 1.0;
        const input = inputs[0];
        const output = outputs[0];
        for (let i = 0; i < input.length; i++) {
            output[i] = input[i] * factor;
        }
    }
};
Ort.env.registerCustomOp(multiplyBy);

Интеграция с WebAssembly и WebGPU

Для моделей, требующих высокопроизводительных вычислений, ORT Web позволяет комбинировать кастомные операторы с WebAssembly или WebGPU бэкендом. В этом случае вычислительная функция может быть реализована с использованием низкоуровневых оптимизаций, при этом интерфейс kernel остаётся стандартным, что обеспечивает совместимость с моделью ONNX.

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

Оглавление