TorchDynamo-based ONNX Exporter

• Обзор
• Зависимости
• Простой пример
• Просмотр ONNX модели с помощью графического интерфейса
• Диагностика проблем с SARIF
• Справочник по API

Обзор

Экспортер ONNX использует движок TorchDynamo для подключения к API оценки фреймов Python и динамического переписывания его байткода в граф FX. Полученный граф FX затем дорабатывается, прежде чем окончательно преобразуется в граф ONNX.

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

Экспортер разработан как модульный и расширяемый. Он состоит из следующих компонентов:

• Экспортер ONNX: Exporter главный класс, который организует процесс экспорта.
• Параметры экспорта ONNX: ExportOptions имеет набор параметров, которые контролируют процесс экспорта.
• Реестр ONNX: OnnxRegistry — это реестр операторов и функций ONNX.
• Экстрактор графа FX: FXGraphExtractor извлекает граф FX из модели PyTorch.
• Режим имитации: ONNXFakeContext — это контекстный менеджер, который активирует режим имитации для моделей большого размера.
• Выход экспорта ONNX: ExportOutput — это выход экспортера, содержащий экспортированный граф ONNX и диагностическую информацию.
• Сериализатор вывода ONNX экспорта: ExportOutputSerializer сериализует экспортированную модель в файл.
• Параметры диагностики ONNX: DiagnosticOptions имеет набор параметров, которые контролируют диагностику, выводимую экспортером.

Зависимости

Экспортер ONNX зависит от дополнительных пакетов Python:

• ONNX
• ONNX Script

Их можно установить через pip:

pip install --upgrade onnx onnxscript

Простой пример

Ниже представлен пример использования API экспортера с простой многослойной перцептронной сетью (MLP):

import torch

class MLPModel(nn.Module):
  def __init__(self):
      super().__init__()
      self.fc0 = nn.Linear(8, 8, bias=True)
      self.fc1 = nn.Linear(8, 4, bias=True)
      self.fc2 = nn.Linear(4, 2, bias=True)
      self.fc3 = nn.Linear(2, 2, bias=True)

  def forward(self, tensor_x: torch.Tensor):
      tensor_x = self.fc0(tensor_x)
      tensor_x = torch.sigmoid(tensor_x)
      tensor_x = self.fc1(tensor_x)
      tensor_x = torch.sigmoid(tensor_x)
      tensor_x = self.fc2(tensor_x)
      tensor_x = torch.sigmoid(tensor_x)
      output = self.fc3(tensor_x)
      return output

model = MLPModel()
tensor_x = torch.rand((97, 8), dtype=torch.float32)
export_output = torch.onnx.dynamo_export(model, tensor_x)

Как показано в коде выше, всё, что вам нужно, это предоставить torch.onnx.dynamo_export() с экземпляром модели и её входными данными. Затем экспортер вернёт экземпляр torch.onnx.ExportOutput, содержащий экспортированный граф ONNX вместе с дополнительной информацией.

Модель в памяти, доступная через export_output.model_proto, — это объект onnx.ModelProto, соответствующий спецификации ONNX IR. Затем ONNX модель может быть сериализована в файл Protobuf с помощью API torch.onnx.ExportOutput.save().

export_output.save("mlp.onnx")

Просмотр ONNX модели с помощью графического интерфейса

Вы можете просмотреть экспортированную модель с помощью Netron.

Обратите внимание, что каждый слой представлен прямоугольной рамкой с иконкой f в правом верхнем углу.

При раскрытии отображается тело функции.

Тело функции — это последовательность операторов ONNX или других функций.

Диагностика проблем с SARIF

Диагностика ONNX выходит за рамки обычных журналов, используя Формат обмена результатами статического анализа (SARIF), чтобы помочь пользователям отлаживать и улучшать свои модели с помощью графического интерфейса, например, Просмотрщика SARIF от Visual Studio Code.

Основные преимущества:

• Диагностика выдаётся в машинно-читаемом формате Формат обмена результатами статического анализа (SARIF).
• Новый более понятный и структурированный способ добавления и отслеживания правил диагностики.
• Служат основой для будущих улучшений, использующих диагностическую информацию.

Правила диагностики ONNX в формате SARIF

• FXE0007:fx-graph-to-onnx
• FXE0008:fx-node-to-onnx
• FXE0010:fx-pass
• FXE0011:отсутствует символическая функция для вызова функции
• FXE0012:неподдерживаемый анализ узла fx
• FXE0013:отладка на уровне операций
• FXE0014:найти совпадение схемы оператора и символической функции
• FXE0015:вставка узла fx для повышения типа
• FXE0016:поиск перегрузок операторов в реестре ONNX

Справочник по API

torch.onnx.dynamo_export(model, /, *model_args, export_options=None, **model_kwargs)

Экспорт модели torch.nn.Module в граф ONNX.

Параметры

• model (Union[Модуль, Callable]) – Модель PyTorch, подлежащая экспорту в ONNX.
• model_args – Позиционные входные данные для model.
• model_kwargs – Именованные входные данные для model.
• export_options (Optional[ExportOptions]) – Параметры, влияющие на экспорт в ONNX.

Возвращает

Представление экспортированной модели ONNX в памяти.

Тип возвращаемого значения

ExportOutput

Пример 1 — Простейший экспорт

class MyModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(2, 2)
    def forward(self, x, bias=None):
        out = self.linear(x)
        out = out + bias
        return out
model = MyModel()
kwargs = {"bias": 3.}
args = (torch.randn(2, 2, 2),)
export_output = torch.onnx.dynamo_export(
    model,
    *args,
    **kwargs).save("my_simple_model.onnx")

Пример 2 — Экспорт с динамическими форматами

# The previous model can be exported with dynamic shapes
export_options = torch.onnx.ExportOptions(dynamic_shapes=True)
export_output = torch.onnx.dynamo_export(
    model,
    *args,
    **kwargs,
    export_options=export_options)
export_output.save("my_dynamic_model.onnx")

Выводом динамических размеров входных данных мы видим, что форма входных данных больше не (2,2,2)

>>> print(export_output.model_proto.graph.input[0])
name: "arg0"
type {
  tensor_type {
    elem_type: 1
    shape {
      dim {
        dim_param: "arg0_dim_0"
      }
      dim {
        dim_param: "arg0_dim_1"
      }
      dim {
        dim_param: "arg0_dim_2"
      }
    }
  }
}

class torch.onnx.ExportOptions(*, dynamic_shapes=None, op_level_debug=None, fake_context=None, onnx_registry=None, diagnostic_options=None)

Параметры, влияющие на экспорт TorchDynamo в ONNX.

Переменные

• dynamic_shapes (Optional[bool]) – Подсказка о форме входных/выходных тензоров. Если None, экспортер определяет наиболее подходящее значение. Если True, все формы входных данных считаются динамическими. Если False, все формы входных данных считаются статическими.
• op_level_debug (Optional[bool]) – Экспортировать ли модель с информацией об отладке на уровне операций
• diagnostic_options (DiagnosticOptions) – Параметры диагностики для экспортера.
• fake_context (Optional[ONNXFakeContext]) – Ложная среда, используемая для символического трассирования.
• onnx_registry (Optional[OnnxRegistry]) – Реестр ONNX, используемый для регистрации операторов ATen в функции ONNX.

torch.onnx.enable_fake_mode()

Включить режим имитации на время выполнения контекста.

Внутренне он создаёт объект-менеджер контекста torch._subclasses.fake_tensor.FakeTensorMode, который преобразует пользовательские входные данные и параметры модели в torch._subclasses.fake_tensor.FakeTensor.

Объект имитации — это torch.Tensor с возможностью выполнения кода PyTorch без фактического выполнения вычислений с тензорами, выделенными на устройстве meta. Поскольку на устройстве не выделяется фактических данных, этот API позволяет экспортировать большие модели без фактической потребности в памяти для их выполнения.

Рекомендуется включать режим имитации при экспорте моделей, слишком больших для размещения в памяти.

Возвращает

Объект ONNXFakeContext, который необходимо передать в dynamo_export() через аргумент ExportOptions.fake_context.

Пример:

# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
>>> import torch
>>> import torch.onnx
>>> class MyModel(torch.nn.Module):  # Dummy model
...     def __init__(self) -> None:
...         super().__init__()
...         self.linear = torch.nn.Linear(2, 2)
...     def forward(self, x):
...         out = self.linear(x)
...         return out
>>> with torch.onnx.enable_fake_mode() as fake_context:
...     my_nn_module = MyModel()
...     arg1 = torch.randn(2, 2, 2)  # positional input 1
>>> export_options = torch.onnx.ExportOptions(fake_context=fake_context)
>>> export_output = torch.onnx.dynamo_export(
...     my_nn_module,
...     arg1,
...     export_options=export_options
... )
>>> # Saving model WITHOUT initializers
>>> export_output.save("my_model_without_initializers.onnx")
>>> # Saving model WITH initializers
>>> export_output.save("my_model_with_initializers.onnx", model_state_dict=MyModel().state_dict())

Предупреждение

Этот API экспериментальный и НЕ совместим с предыдущими версиями.

class torch.onnx.ExportOutput(model_proto, input_adapter, output_adapter, diagnostic_context, *, fake_context=None, export_exception=None)

Представление модели PyTorch, экспортированной в ONNX, в оперативной памяти.

adapt_torch_inputs_to_onnx(*model_args, **model_kwargs) [source]

Преобразует входные данные модели PyTorch в формат входных данных экспортированной модели ONNX.

Из-за различий в проектировании формат входных/выходных данных модели PyTorch и экспортированной модели ONNX часто различаются. Например, None допускается для модели PyTorch, но не поддерживается ONNX. Вложенные структуры тензоров допускаются для модели PyTorch, но ONNX поддерживает только плоские тензоры и т.д.

Фактические шаги адаптации связаны с каждым отдельным экспортом. Они зависят от модели PyTorch, конкретного набора model_args и model_kwargs, используемых для экспорта, и параметров экспорта.

Этот метод воспроизводит шаги адаптации, записанные во время экспорта.

Параметры

• model_args – Входные данные модели PyTorch.
• model_kwargs – Параметрические входные данные модели PyTorch.

Возвращает

Последовательность тензоров, преобразованных из входных данных модели PyTorch.

Тип возвращаемого значения

Sequence[Union[Tensor, int, float, bool]]

Пример:

# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
>>> import torch
>>> import torch.onnx
>>> from typing import Dict, Tuple
>>> def func_with_nested_input_structure(
...     x_dict: Dict[str, torch.Tensor],
...     y_tuple: Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
... ):
...     if "a" in x_dict:
...         x = x_dict["a"]
...     elif "b" in x_dict:
...         x = x_dict["b"]
...     else:
...         x = torch.randn(3)
...
...     y1, (y2, y3) = y_tuple
...
...     return x + y1 + y2 + y3
>>> x_dict = {"a": torch.tensor(1.)}
>>> y_tuple = (torch.tensor(2.), (torch.tensor(3.), torch.tensor(4.)))
>>> export_output = torch.onnx.dynamo_export(func_with_nested_input_structure, x_dict, y_tuple)
>>> print(x_dict, y_tuple)
{'a': tensor(1.)} (tensor(2.), (tensor(3.), tensor(4.)))
>>> print(export_output.adapt_torch_inputs_to_onnx(x_dict, y_tuple))
(tensor(1.), tensor(2.), tensor(3.), tensor(4.))

Предупреждение

Этот API является экспериментальным и НЕ обратной совместимым.

adapt_torch_outputs_to_onnx(model_outputs) [source]

Преобразует выходные данные модели PyTorch в формат выходных данных экспортированной модели ONNX.

Из-за различий в проектировании формат входных/выходных данных модели PyTorch и экспортированной модели ONNX часто различаются. Например, None допускается для модели PyTorch, но не поддерживается ONNX. Вложенные структуры тензоров допускаются для модели PyTorch, но ONNX поддерживает только плоские тензоры и т.д.

Фактические шаги адаптации связаны с каждым отдельным экспортом. Они зависят от модели PyTorch, конкретного набора model_args и model_kwargs, используемых для экспорта, и параметров экспорта.

Этот метод воспроизводит шаги адаптации, записанные во время экспорта.

Параметры

model_outputs (Any) – Выходные данные модели PyTorch.

Возвращает

Выходные данные модели PyTorch в формате выходных данных экспортированной модели ONNX.

Тип возвращаемого значения

Sequence[Union[Tensor, int, float, bool]]

Пример:

# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
>>> import torch
>>> import torch.onnx
>>> def func_returning_tuples(x, y, z):
...     x = x + y
...     y = y + z
...     z = x + y
...     return (x, (y, z))
>>> x = torch.tensor(1.)
>>> y = torch.tensor(2.)
>>> z = torch.tensor(3.)
>>> export_output = torch.onnx.dynamo_export(func_returning_tuples, x, y, z)
>>> pt_output = func_returning_tuples(x, y, z)
>>> print(pt_output)
(tensor(3.), (tensor(5.), tensor(8.)))
>>> print(export_output.adapt_torch_outputs_to_onnx(pt_output))
[tensor(3.), tensor(5.), tensor(8.)]

Предупреждение

Этот API является экспериментальным и НЕ обратной совместимым.

property diagnostic_context: diagnostics.DiagnosticContext

Контекст диагностики, связанный с экспортом.

property fake_context: Optional[ONNXFakeContext]

Фейковый контекст, связанный с экспортом.

property model_proto: onnx.ModelProto

Экспортированная модель ONNX в формате onnx.ModelProto.

save(destination, *, model_state_dict=None, serializer=None) [source]

Сохраняет модель ONNX в оперативной памяти в destination с указанными serializer.

Параметры

• destination (Union[str, BufferedIOBase]) – Место сохранения модели ONNX. Может быть строкой или объектом типа «поток». При использовании с model_state_dict, это должна быть строка с полным путем к месту назначения. В этом случае помимо сохранения модели ONNX, будет создана папка с суффиксом «_initializers» (без расширения) для хранения каждого инициализатора модели ONNX в отдельном файле. Например, если место назначения — «/path/model.onnx», инициализаторы будут сохранены в папке «/path/model_initializers/».
• model_state_dict (Optional[Union[Dict[str, Any], str]]) – state_dict модели PyTorch, содержащий все веса. Может быть словарем, возвращаемым model.state_dict(), или строкой с именем файла. Требуется при использовании enable_fake_mode(), но реальные инициализаторы требуются в графе ONNX. Может быть либо строкой с путем к контрольной точке, либо словарем с фактическим состоянием модели.
• serializer (Optional[ExportOutputSerializer]) – Использовать сериализатор. Если не указано, модель будет сериализована как Protobuf.

save_diagnostics(destination) [source]

Сохраняет диагностику экспорта в журнал SARIF по указанному пути.

Параметры

destination (str) – Место сохранения журнала диагностики SARIF. Должен иметь расширение .sarif.

Исключения

ValueError – Если путь к месту назначения не заканчивается расширением .sarif.

class torch.onnx.ExportOutputSerializer(*args, **kwargs)

Протокол для сериализации графа ONNX в определённый формат (например, Protobuf). Обратите внимание, что это продвинутый сценарий использования.

serialize(export_output, destination) [source]

Метод протокола, который должен быть реализован для сериализации.

Параметры

• export_output (ExportOutput) – Представляет модель ONNX, экспортированную в оперативной памяти.
• destination (BufferedIOBase) – Бинарный поток ввода-вывода или предварительно выделенный буфер, в который следует записать сериализованную модель.

Пример

Простой сериализатор, записывающий экспортированный onnx.ModelProto в формате Protobuf в destination.

# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
>>> import io
>>> import torch
>>> import torch.onnx
>>> class MyModel(torch.nn.Module):  # Dummy model
...     def __init__(self) -> None:
...         super().__init__()
...         self.linear = torch.nn.Linear(2, 2)
...     def forward(self, x):
...         out = self.linear(x)
...         return out
>>> class ProtobufExportOutputSerializer:
...     def serialize(
...         self, export_output: torch.onnx.ExportOutput, destination: io.BufferedIOBase
...     ) -> None:
...         destination.write(export_output.model_proto.SerializeToString())
>>> model = MyModel()
>>> arg1 = torch.randn(2, 2, 2)  # positional input 1
>>> torch.onnx.dynamo_export(model, arg1).save(
...     destination="exported_model.onnx",
...     serializer=ProtobufExportOutputSerializer(),
... )

class torch.onnx.OnnxExporterError(export_output, message)

Возникает при ошибке экспортера ONNX.

Это исключение выбрасывается при ошибке во время процесса экспорта ONNX. Оно инкапсулирует объект ExportOutput, созданный до момента возникновения ошибки, позволяя получить частичные результаты экспорта и связанную метаданные.

class torch.onnx.OnnxRegistry

Реестр функций ONNX.

Реестр поддерживает сопоставление квалифицированных имен с символическими функциями в рамках фиксированной версии opset. Он поддерживает регистрацию пользовательских функций onnx-script и для диспетчера, чтобы перенаправлять вызовы на соответствующую функцию.

get_op_functions(namespace, op_name, overload=None) [source]

Возвращает список ONNXFunctions для данного оператора: torch.ops.<namespace>.<op_name>.<overload>.

Список упорядочен по времени регистрации. Пользовательские операторы должны быть во второй половине списка.

Параметры

• namespace (строка) – Имя пространства имен оператора.
• op_name (строка) – Имя оператора.
• overload (Необязательный[строка]) – Перегрузка оператора. Если это стандартная перегрузка, оставьте None.

Возвращает

Список ONNXFunctions, соответствующих заданному имени, или None, если имя отсутствует в реестре.

Тип возвращаемого значения

Необязательный[Список[ONNXFunction]]

is_registered_op(namespace, op_name, overload=None) [source]

Возвращает, зарегистрирован ли данный оператор: torch.ops.<namespace>.<op_name>.<overload>.

Параметры

• namespace (строка) – Имя пространства имен оператора.
• op_name (строка) – Имя оператора.
• overload (Необязательный[строка]) – Перегрузка оператора. Если это стандартная перегрузка, оставьте None.

Возвращает

True, если оператор зарегистрирован, иначе False.

Тип возвращаемого значения

bool

property opset_version: int

Версия ONNX opset, на которую должен ориентироваться экспортер. По умолчанию это последняя поддерживаемая версия ONNX opset: 18. Значение по умолчанию будет увеличиваться со временем по мере развития ONNX.

register_op(function, namespace, op_name, overload=None, is_complex=False) [source]

Регистрирует пользовательский оператор: torch.ops.<namespace>.<op_name>.<overload>.

Параметры

• function (Объединение[onnxscript.OnnxFunction, onnxscript.TracedOnnxFunction]) – Функция onnx-script для регистрации.
• namespace (строка) – Имя пространства имен оператора.
• op_name (строка) – Имя оператора.
• overload (Необязательный[строка]) – Перегрузка оператора. Если это стандартная перегрузка, оставьте None.
• is_complex (bool) – Является ли функция функцией, обрабатывающей комплексные входные данные.

Исключения

ValueError – Если имя не имеет вида «namespace::op».

class torch.onnx.DiagnosticOptions(verbosity_level=20, warnings_as_errors=False)

Параметры для диагностического контекста.

Переменные

• verbosity_level (целое число) – Устанавливает количество информации, регистрируемой для каждой диагностики, эквивалентно «уровню» в модуле Python logging.
• warnings_as_errors (bool) – Когда True, предупреждения-диагностики рассматриваются как ошибки-диагностики.