Обрезка графа и удаление лишних узлов

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

Принципы обрезки графа

ONNX-модель представляет собой направленный ациклический граф (DAG), где узлы — это операции (операторы ONNX), а ребра — тензоры, передающиеся между операциями. В процессе подготовки модели к веб-выполнению часто встречаются:

  • Вспомогательные узлы, которые не используются для вычисления конечного выхода.
  • Узлы, связанные с обучением, например Dropout или BatchNormalization, которые могут быть упрощены или удалены.
  • Неиспользуемые промежуточные вычисления, созданные при экспортировании модели из фреймворка.

Удаление этих узлов позволяет сократить граф до минимально необходимого набора операций.

Ключевые шаги:

  1. Анализ зависимостей узлов

    • Каждому узлу присваивается список входов и выходов.
    • Строится карта достижимости: от каждого узла отмечается, какие выходные данные он влияет.

  2. Определение корневых узлов

    • Корневые узлы — это узлы, результаты которых используются для вычисления целевых выходов модели (model.graph.output).
    • Все узлы, не влияющие на эти выходы, считаются кандидатом на удаление.

  3. Рекурсивное удаление

    • Проход по графу начиная с выходных узлов.
    • Все узлы, до которых нет пути к выходу, удаляются.

Инструменты ONNX Runtime Web для обрезки

ONNX Runtime Web не предоставляет встроенного интерфейса для прямого редактирования графа на уровне JavaScript, но поддерживает работу с оптимизированными моделями и оптимизацию при загрузке через InferenceSession:

import * as ort from 'onnxruntime-web';

const sessionOptions = {
  executionProviders: ['wasm'],
  graphOptimizationLevel: 'all' // full optimization, включая обрезку графа
};

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

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

  • graphOptimizationLevel

    • disabled — без оптимизации
    • basic — базовые оптимизации, удаление констант и простой dead code
    • extended — расширенные оптимизации, включая слияние узлов
    • all — полная оптимизация, включая обрезку неиспользуемых узлов

  • executionProviders

    • Определяет, где будет выполняться модель (wasm, webgl), что важно при оптимизации для браузера.

Удаление лишних узлов вручную

Для сложных случаев иногда необходимо вручную модифицировать ONNX-граф до передачи в Web. Процесс включает:

  1. Загрузка модели с помощью ONNX.js или Python ONNX

    • Обычно это делается с помощью Python-пакета onnx.
    • Модель читается в onnx.ModelProto, где можно получить доступ к graph.node.

  2. Анализ каждого узла

    import onnx

mo del = onnx.load("model.onnx")
graph = model.graph

# Сбор всех используемых тензоров
used_tensors = set(output.name for output in graph.output)
for node in reversed(graph.node):
    if any(out in used_tensors for out in node.output):
        used_tensors.update(node.input)
    else:
        graph.node.remove(node)
onnx.save(model, "pruned_model.onnx")

    Этот пример удаляет все узлы, которые не влияют на выходы модели.

  3. Экспорт оптимизированной модели в Web

    • После удаления лишних узлов модель становится легче и быстрее загружается в браузере.
    • Уменьшение количества операций также сокращает время исполнения при использовании ort.InferenceSession.

Практические рекомендации

  • Минимизировать промежуточные вычисления: удалять узлы типа Identity, Dropout и ненужные константы.
  • Использовать graphOptimizationLevel: 'all' для максимальной автоматической обрезки графа.
  • Проверять соответствие выходов: после удаления узлов важно убедиться, что значения конечных тензоров остаются корректными.
  • Интегрировать обрезку в CI/CD: модели можно автоматически оптимизировать перед публикацией веб-приложения, снижая нагрузку на клиент.

Влияние на производительность

Удаление лишних узлов:

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

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

Оглавление