Динамические графы и условные операторы в ONNX

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

В ONNX динамический граф реализуется через поддержку тензоров с динамическими размерами, условных ветвлений и циклов. Динамическая размерность задается как символическая, например batch_size или sequence_length, что позволяет модели принимать на вход последовательности различной длины или пакеты различного объёма без необходимости пересоздания графа.

Пример динамического входа в ONNX Runtime Web:

import * as ort from 'onnxruntime-web';

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

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

const feeds = { input: inputTensor };

const results = await session.run(feeds);
console.log(results.output.data);

Ключевой момент — возможность передавать тензоры с различной формой, если модель была экспортирована с динамическими размерностями.

Условные операторы

ONNX поддерживает несколько операторов для реализации условной логики: If, Loop, а также операторы сравнения и булевой логики (Greater, Less, Equal, And, Or, Not).

Оператор If позволяет выполнять различные подграфы в зависимости от логического условия:

// Псевдографика условного оператора в ONNX
// input_condition -> If
//    then_branch: выполняется, если условие true
//    else_branch: выполняется, если условие false

Для ORT Web условный оператор работает полностью на стороне клиента. Условие подается как тензор типа bool:

const condition = new ort.Tensor('bool', new Uint8Array([1]), [1]);

const feeds = { condition_input: condition, input_data: inputTensor };

const results = await session.run(feeds);

Внутри модели If выбирает соответствующий подграф. Это позволяет создавать модели с ветвлением, например, для разных типов обработки данных, без изменения структуры модели после экспорта в ONNX.

Циклы и повторяющиеся вычисления

Для динамических последовательностей и рекуррентных операций используется оператор Loop. Он позволяет выполнять вычисления несколько раз, пока не выполнится условие остановки, определяемое логическим тензором.

Основные элементы Loop:

  • M — максимальное количество итераций.
  • cond — булевый тензор, определяющий, продолжать цикл или нет.
  • loop-carried dependencies — значения, передаваемые из одной итерации в другую.

Пример применения:

// session уже создан
const maxIters = new ort.Tensor('int64', BigInt64Array.from([10]), [1]);
const cond = new ort.Tensor('bool', new Uint8Array([1]), [1]);
const loopVars = { iter_count: maxIters, continue_loop: cond };

const results = await session.run(loopVars);

Через Loop можно реализовать сложные алгоритмы с динамическими условиями, например, рекуррентные сети, условные итерации или генерацию последовательностей переменной длины.

Совмещение динамических графов и условной логики

Модели с динамическими графами и условными операторами позволяют создавать адаптивные вычислительные схемы, которые:

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

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

Особенности оптимизации в ONNX Runtime Web

  1. JIT-компиляция под WebAssembly и WebGPU. Динамические ветвления и циклы обрабатываются эффективно, даже при изменении формы входных тензоров.
  2. Кэширование промежуточных тензоров для повторного использования в циклах.
  3. Минимизация трансформаций между CPU и GPU, что ускоряет обработку динамических графов.

Практические советы

  • Всегда указывать символические размеры при экспорте модели, если планируется использовать динамический ввод.
  • При работе с If и Loop в браузере проверять корректность типов булевых и целочисленных тензоров.
  • Для больших моделей динамических графов следует учитывать ограничения памяти браузера: оптимизация размера батча и промежуточных тензоров критична.

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

Оглавление