Spec-Zone.ru › PyTorch 2.14

Экспортёр ONNX на основе torch.export

Создано: 10 июн. 2025 | Последнее обновление: 30 янв. 2026

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

Обзор

Для создания трассированного графа, представляющего только вычисления с тензорами в функции, в режиме опережающей компиляции (Ahead-of-Time, AOT) используется механизм torch.export. Полученный трассированный граф (1) содержит нормализованные операторы из функционального набора операторов ATen (а также любые заданные пользователем пользовательские операторы), (2) не содержит потока управления и структур данных Python (за некоторыми исключениями) и (3) записывает набор ограничений на формы, необходимых для подтверждения корректности такой нормализации и устранения потока управления для будущих входных данных, после чего преобразуется в граф ONNX.

Кроме того, в процессе экспорта значительно сокращается использование памяти.

Зависимости

Для работы экспортёра ONNX требуются дополнительные пакеты Python:

  • ONNX
  • ONNX Script

Их можно установить с помощью pip:

  pip install --upgrade onnx onnxscript

Затем onnxruntime можно использовать для выполнения модели на самых разных процессорах.

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

Ниже показана работа API экспортёра на простом многослойном перцептроне (MLP):

import torch
import torch.nn as nn

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)
      self.fc_combined = nn.Linear(8 + 8 + 8, 8, bias=True)  # Combine all inputs

  def forward(self, tensor_x: torch.Tensor, input_dict: dict, input_list: list):
      """
      Forward method that requires all inputs:
      - tensor_x: A direct tensor input.
      - input_dict: A dictionary containing the tensor under the key 'tensor_x'.
      - input_list: A list where the first element is the tensor.
      """
      # Extract tensors from inputs
      dict_tensor = input_dict['tensor_x']
      list_tensor = input_list[0]

      # Combine all inputs into a single tensor
      combined_tensor = torch.cat([tensor_x, dict_tensor, list_tensor], dim=1)

      # Process the combined tensor through the layers
      combined_tensor = self.fc_combined(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      combined_tensor = self.fc0(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      combined_tensor = self.fc1(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      combined_tensor = self.fc2(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      output = self.fc3(combined_tensor)
      return output

model = MLPModel()

# Example inputs
tensor_input = torch.rand((97, 8), dtype=torch.float32)
dict_input = {'tensor_x': torch.rand((97, 8), dtype=torch.float32)}
list_input = [torch.rand((97, 8), dtype=torch.float32)]

# The input_names and output_names are used to identify the inputs and outputs of the ONNX model
input_names = ['tensor_input', 'tensor_x', 'list_input_index_0']
output_names = ['output']

# Exporting the model with all required inputs
onnx_program = torch.onnx.export(model,(tensor_input, dict_input, list_input), dynamic_shapes=({0: "batch_size"},{"tensor_x": {0: "batch_size"}},[{0: "batch_size"}]), input_names=input_names, output_names=output_names, dynamo=True,)

# Check the exported ONNX model is dynamic
assert onnx_program.model.graph.inputs[0].shape == ("batch_size", 8)
assert onnx_program.model.graph.inputs[1].shape == ("batch_size", 8)
assert onnx_program.model.graph.inputs[2].shape == ("batch_size", 8)

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

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

  onnx_program.save("mlp.onnx")

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

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

MLP model as viewed using Netron

Если преобразование завершается с ошибкой

Функцию torch.onnx.export() следует вызвать ещё раз с параметром report=True. Для помощи в устранении проблемы пользователем создаётся отчёт в формате Markdown.

Метаданные

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

К каждому узлу ONNX добавляются следующие поля метаданных:

  • namespace

    Строка, представляющая иерархическое пространство имён узла и состоящая из стека вызовов модулей и методов.

    Пример: __main__.SimpleAddModel/add: aten.add.Tensor

  • pkg.torch.onnx.class_hierarchy

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

    Пример: ['__main__.SimpleAddModel', 'aten.add.Tensor']

  • pkg.torch.onnx.fx_node

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

    Пример: %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%tensor_x, %input_dict_tensor_x, %input_list_0], 1), kwargs = {})

  • pkg.torch.onnx.name_scopes

    Список областей имён (методов), представляющих путь к этому узлу в модели PyTorch.

    Пример: ['', 'add']

  • pkg.torch.onnx.stack_trace

    Стек вызовов из исходного кода, в котором был создан этот узел, если он доступен.

    Пример:

    File "simpleadd.py", line 7, in forward
    return torch.add(x, y)

Эти поля метаданных хранятся в атрибуте metadata_props каждого узла ONNX; их можно просмотреть с помощью Netron или программно.

Для графа ONNX в целом доступны следующие metadata_props:

  • pkg.torch.export.ExportedProgram.graph_signature

    Это свойство содержит строковое представление graph_signature исходной программы ExportedProgram в PyTorch. Сигнатура графа описывает структуру входных и выходных данных модели и их соответствие графу ONNX. Входные данные определяются как объекты InputSpec, которые включают тип входных данных (например, InputKind.PARAMETER для параметров, InputKind.USER_INPUT для входных данных, заданных пользователем), имя аргумента, целевой объект (которым может быть определённый узел модели) и признак постоянства входных данных. Выходные данные определяются как объекты OutputSpec, указывающие тип выходных данных (например, OutputKind.USER_OUTPUT) и имя аргумента.

    Подробнее о сигнатуре графа см. в документации torch.export.

  • pkg.torch.export.ExportedProgram.range_constraints

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

    Пример: s0: VR[2, int_oo], указывающий, что размер входного тензора должен быть не меньше 2.

    Подробнее об ограничениях диапазонов см. в документации torch.export.

Для каждого входного значения графа ONNX могут присутствовать следующие свойства метаданных:

  • pkg.torch.export.graph_signature.InputSpec.kind

    Тип входных данных, определённый перечислением InputKind в PyTorch.

    Примеры значений:

    • “USER_INPUT”: входные данные модели, предоставленные пользователем.
    • “PARAMETER”: параметр модели (например, вес).
    • “BUFFER”: буфер модели (например, скользящее среднее в BatchNorm).
    • “CONSTANT_TENSOR”: аргумент — тензор-константа.
    • “CUSTOM_OBJ”: входные данные — пользовательский объект.
    • “TOKEN”: входной токен.

  • pkg.torch.export.graph_signature.InputSpec.persistent

    Указывает, являются ли входные данные постоянными (то есть должны ли они сохраняться как часть состояния модели).

    Примеры значений:

    • “True”
    • “False”

Для каждого выходного значения графа ONNX могут присутствовать следующие свойства метаданных:

  • pkg.torch.export.graph_signature.OutputSpec.kind

    Тип выходных данных, определённый перечислением OutputKind в PyTorch.

    Примеры значений:

    • “USER_OUTPUT”: выходные данные, видимые пользователю.
    • “LOSS_OUTPUT”: выходное значение функции потерь.
    • “BUFFER_MUTATION”: указывает, что буфер был изменён.
    • “GRADIENT_TO_PARAMETER”: выходные данные с градиентом параметра.
    • “GRADIENT_TO_USER_INPUT”: выходные данные с градиентом входных данных пользователя.
    • “USER_INPUT_MUTATION”: указывает, что входные данные пользователя были изменены.
    • “TOKEN”: выходной токен.

Каждое инициализированное значение, входные и выходные данные имеют следующие метаданные:

  • pkg.torch.onnx.original_node_name

    Исходное имя узла графа FX в PyTorch, создавшего это значение, если значение было переименовано. Это помогает проследить инициализаторы до их источника в исходной модели.

    Пример: fc1.weight

Справочник API

torch.onnx.export(model, args=(), f=None, *, kwargs=None, verbose=None, input_names=None, output_names=None, opset_version=None, dynamo=True, external_data=True, dynamic_shapes=None, custom_translation_table=None, report=False, optimize=True, verify=False, profile=False, dump_exported_program=False, artifacts_dir='.', export_params=True, keep_initializers_as_inputs=False, dynamic_axes=None, training=<TrainingMode.EVAL: 0>, operator_export_type=<OperatorExportTypes.ONNX: 0>, do_constant_folding=True, custom_opsets=None, export_modules_as_functions=False, autograd_inlining=True) [исходный код]

Экспортирует модель в формат ONNX.

Установка dynamo=True включает новую логику экспорта ONNX, основанную на torch.export.ExportedProgram и более современном наборе правил преобразования. Это рекомендуемый способ экспорта моделей в ONNX, используемый по умолчанию.

Когда dynamo=True:

Экспортёр использует следующие стратегии, чтобы получить ExportedProgram для преобразования в ONNX.

  1. Если модель уже является ExportedProgram, она будет использоваться без изменений.
  2. Использовать torch.export.export() и установить strict=False.
  3. Использовать torch.export.export() и установить strict=True.
Параметры:
  • model (torch.nn.Module | torch.export.ExportedProgram | torch.jit.ScriptModule | torch.jit.ScriptFunction) – Модель для экспорта.
  • args (tuple[Any, ...]) – Примеры позиционных входных данных. Все аргументы, не являющиеся Tensor, будут жестко заданы в экспортированной модели; аргументы Tensor станут входными данными экспортированной модели в том порядке, в котором они указаны в кортеже.
  • f (str | os.PathLike | None) – Путь к файлу выходной модели ONNX. Например, «model.onnx». Этот аргумент сохранён для обратной совместимости. Рекомендуется не указывать его (None) и вместо этого использовать возвращённый torch.onnx.ONNXProgram для сериализации модели в файл.
  • kwargs (dict[str, Any] | None) – Необязательные примеры именованных входных данных.
  • verbose (bool | None) – Включить ли подробное журналирование.
  • input_names (Sequence[str] | None) – имена, назначаемые входным узлам графа, по порядку.
  • output_names (Sequence[str] | None) – имена, назначаемые выходным узлам графа, по порядку. Это только метки, они не влияют на порядок выходных данных. Если модель возвращает словарь, выходные данные разворачиваются в порядке итерации по словарю независимо от указанных здесь имён.
  • opset_version (int | None) – Версия набора операторов по умолчанию (ai.onnx), на которую нужно ориентироваться. Укажите opset_version в соответствии с версиями набора операторов, поддерживаемыми целевой средой выполнения или компилятором, в которых будет использоваться экспортированная модель. Оставьте значение по умолчанию (None), чтобы использовать рекомендуемую версию, либо обратитесь к документации по операторам ONNX для получения дополнительной информации.
  • dynamo (bool) – Экспортировать ли модель с помощью ExportedProgram на основе torch.export вместо TorchScript.
  • external_data (bool) – Сохранять ли веса модели во внешнем файле данных. Это требуется для моделей с большими весами, превышающими ограничение размера файла ONNX (2 ГБ). Если значение False, веса сохраняются в файле ONNX вместе с архитектурой модели.
  • dynamic_shapes (dict[str, Any] | tuple[Any, ...] | list[Any] | None) – Словарь или кортеж динамических форм входных данных модели. Дополнительные сведения см. в torch.export.export(). Используется (и предпочтителен) только при значении dynamo=True. Обратите внимание: dynamic_shapes предназначен для экспорта модели с dynamo=True, тогда как dynamic_axes используется при dynamo=False.
  • custom_translation_table (dict[Callable, Callable] | None) – Словарь пользовательских декомпозиций операторов модели. Ключом словаря должна быть целевая вызываемая сущность в узле fx (например, torch.ops.aten.stft.default), а значением — функция, строящая этот граф с помощью ONNX Script. Этот параметр действителен только при dynamo=True.
  • report (bool) – Создавать ли отчёт в формате Markdown о процессе экспорта. Этот параметр действителен только при dynamo=True.
  • optimize (bool) – Оптимизировать ли экспортированную модель. Этот параметр действителен только при dynamo=True. Значение по умолчанию — True.
  • verify (bool) – Проверять ли экспортированную модель с помощью ONNX Runtime. Этот параметр действителен только при dynamo=True.
  • profile (bool) – Профилировать ли процесс экспорта. Этот параметр действителен только при dynamo=True.
  • dump_exported_program (bool) – Сохранять ли torch.export.ExportedProgram в файл. Это полезно для отладки экспортёра. Этот параметр действителен только при dynamo=True.
  • artifacts_dir (str | os.PathLike) – Каталог для сохранения отладочных артефактов, например отчёта и сериализованной экспортированной программы. Этот параметр действителен только при dynamo=True.

  • export_params (bool) –

    Если указан ``f``: если значение false, параметры (веса) не будут экспортированы.

    Также можно не указывать этот параметр и использовать возвращённый torch.onnx.ONNXProgram, чтобы управлять обработкой инициализаторов при сериализации модели.

  • keep_initializers_as_inputs (bool) –

    Если указан ``f``: если значение True, все инициализаторы (обычно соответствующие весам модели) в экспортированном графе также будут добавлены в граф как входные данные. Если значение False, инициализаторы не добавляются в граф как входные данные, и входными данными становятся только пользовательские входы.

    Установите True, если планируете передавать веса модели во время выполнения. Установите False, если веса статические, чтобы обеспечить более эффективную оптимизацию (например, свёртку констант) бэкендами и средами выполнения.

    Также можно не указывать этот параметр и использовать возвращённый torch.onnx.ONNXProgram, чтобы управлять обработкой инициализаторов при сериализации модели.

  • dynamic_axes (Mapping[str, Mapping[int, str]] | Mapping[str, Sequence[int]] | None) –

    Устарело: при dynamo=True рекомендуется указывать dynamic_shapes.

    По умолчанию формы всех входных и выходных тензоров экспортированной модели будут в точности соответствовать формам, заданным в args. Чтобы указать динамические оси тензоров (то есть известные только во время выполнения), задайте dynamic_axes в виде словаря со следующей схемой:

    • KEY (str): an input or output name. Each name must also be provided in input_names or

      output_names.

    • ЗНАЧЕНИЕ (dict или list): если это dict, ключами являются индексы осей, а значениями — имена осей. Если это

      list, каждый элемент является индексом оси.

    Например:

    class SumModule(torch.nn.Module):
    def forward(self, x):
        return torch.sum(x, dim=1)


torch.onnx.export(
    SumModule(),
    (torch.ones(2, 2),),
    "onnx.pb",
    input_names=["x"],
    output_names=["sum"],
)

    Даёт:

    input {
  name: "x"
  ...
      shape {
        dim {
          dim_value: 2  # axis 0
        }
        dim {
          dim_value: 2  # axis 1
...
output {
  name: "sum"
  ...
      shape {
        dim {
          dim_value: 2  # axis 0
...

    Тогда как:

    torch.onnx.export(
    SumModule(),
    (torch.ones(2, 2),),
    "onnx.pb",
    input_names=["x"],
    output_names=["sum"],
    dynamic_axes={
        # dict value: manually named axes
        "x": {0: "my_custom_axis_name"},
        # list value: automatic names
        "sum": [0],
    },
)

    Даёт:

    input {
  name: "x"
  ...
      shape {
        dim {
          dim_param: "my_custom_axis_name"  # axis 0
        }
        dim {
          dim_value: 2  # axis 1
...
output {
  name: "sum"
  ...
      shape {
        dim {
          dim_param: "sum_dynamic_axes_1"  # axis 0
...
  • training (_C_onnx.TrainingMode) – Устаревший параметр. Вместо этого задайте режим обучения модели перед экспортом.
  • operator_export_type (_C_onnx.OperatorExportTypes) – Устаревший параметр. Поддерживается только ONNX.
  • do_constant_folding (bool) – Устаревший параметр.
  • custom_opsets (Mapping[str, int] | None) – Устаревший параметр.
  • export_modules_as_functions (bool | Collection[type[torch.nn.Module]]) – Устаревший параметр.
  • autograd_inlining (bool) – Устаревший параметр.
Возвращает:

torch.onnx.ONNXProgram, если dynamo имеет значение True, иначе None.

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

ONNXProgram | None

Изменено в версии 2.6: Параметр training теперь устарел. Вместо этого задайте режим обучения модели перед экспортом. Параметр operator_export_type теперь устарел. Поддерживается только ONNX. Параметр do_constant_folding теперь устарел. Он всегда включён. Параметр export_modules_as_functions теперь устарел. Параметр autograd_inlining теперь устарел.

Изменено в версии 2.7: Теперь по умолчанию имеет значение True: optimize.

Изменено в версии 2.9: Теперь по умолчанию имеет значение True: dynamo.

Изменено в версии 2.11: Параметр fallback удалён.

class torch.onnx.ONNXProgram(model, exported_program)

Класс, представляющий программу ONNX, которую можно вызвать с тензорами torch.

Переменные:
  • model – Модель ONNX в виде объекта модели ONNX IR.
  • exported_program – Экспортированная программа, создавшая модель ONNX.
apply_weights(state_dict) [исходный код]

Применяет к модели ONNX веса из указанного словаря состояний.

Используйте этот метод для замены FakeTensors или других весов.

Параметры:

state_dict (dict[str, Tensor]) – Словарь состояний, содержащий веса для применения к модели ONNX.

call_reference(*args, **kwargs) [исходный код]

Запускает модель ONNX с помощью эталонного бэкенда.

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

Sequence[Tensor]

compute_values(value_names, args=(), kwargs=None) [исходный код]

Вычисляет значения указанных имён в модели ONNX.

Этот метод вычисляет значения указанных имён в модели ONNX. Значения возвращаются в виде словаря, сопоставляющего имена с тензорами.

Параметры:

value_names (Sequence[str]) – Имена вычисляемых значений.

Возвращает:

Словарь, сопоставляющий имена с тензорами.

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

Sequence[Tensor]

initialize_inference_session(initializer=<function _ort_session_initializer>) [исходный код]

Инициализирует сеанс выполнения ONNX Runtime.

Параметры:

initializer (Callable[[str | bytes], ort.InferenceSession]) – Функция для инициализации сеанса выполнения ONNX Runtime с указанной моделью. По умолчанию используется функция _ort_session_initializer().

property model_proto: ModelProto

Возвращает объект ONNX ModelProto.

optimize() [исходный код]

Оптимизирует модель ONNX.

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

release() [исходный код]

Освобождает сеанс выполнения.

Вызовите этот метод, чтобы освободить ресурсы, используемые сеансом выполнения.

rename_axes(rename_mapping) [исходный код]

Переименовывает оси модели в соответствии с указанным сопоставлением имён.

Пример:

batch = onnx_program.model.graph.inputs[0].shape[0]
seq_len = onnx_program.model.graph.inputs[0].shape[2]
rename_mapping = {
    batch: "batch",
    seq_len: "seq_len",
}
onnx_program.rename_axes(rename_mapping)
Параметры:

rename_mapping (dict[str | SymbolicDim, str]) –

Словарь, сопоставляющий старые оси с новыми именами осей. В качестве ключей можно использовать:

  • Строковые имена осей (например, «s1», «s2»)
  • Объекты SymbolicDim, полученные из модели (например, onnx_program.model.graph.inputs[0].shape[0])

Значения должны быть строками, представляющими новые имена осей.

save(destination, *, include_initializers=True, keep_initializers_as_inputs=False, external_data=None) [исходный код]

Сохраняет модель ONNX в указанное место назначения.

Если external_data — True или размер модели превышает 2 ГБ, веса сохраняются как внешние данные в отдельном файле.

Способы сериализации инициализаторов (весов модели):

  • include_initializers=True, keep_initializers_as_inputs=False (по умолчанию): инициализаторы включаются в сохранённую модель.
  • include_initializers=True, keep_initializers_as_inputs=True: инициализаторы включаются в сохранённую модель и остаются входными данными модели. Выберите этот вариант, если хотите иметь возможность переопределять веса модели во время инференса.
  • include_initializers=False, keep_initializers_as_inputs=False: инициализаторы не включаются в сохранённую модель и не указываются как входные данные модели. Выберите этот вариант, если хотите добавить инициализаторы в модель ONNX отдельным этапом постобработки.
  • include_initializers=False, keep_initializers_as_inputs=True: инициализаторы не включаются в сохранённую модель, но указываются как входные данные модели. Выберите этот вариант, если хотите передавать инициализаторы во время инференса и уменьшить размер сохранённой модели.
Параметры:
  • destination (str | PathLike) – Путь для сохранения модели ONNX.
  • include_initializers (bool) – Включать ли инициализаторы в сохранённую модель.
  • keep_initializers_as_inputs (bool) – Оставлять ли инициализаторы в качестве входных данных сохранённой модели. Если True, инициализаторы добавляются в модель как входные данные, а значит, их можно переопределить, передав инициализаторы в качестве входных данных модели.
  • external_data (bool | None) – Сохранять ли веса как внешние данные в отдельном файле.
Вызывает исключение:

TypeError – Если external_data — True, а destination не является путём к файлу.

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

Абстрактный интерфейс для модулей, экспортируемых в ONNX.

Унаследуйте этот класс и реализуйте определённые абстрактные методы, чтобы создать модуль, который можно экспортировать в формат ONNX.

Пример:

class Model(torch.nn.Module):
    def forward(self, x):
        return x * 2


class MyExportableModule(torch.onnx.ExportableModule):
    def __init__(self):
        super().__init__()
        self.model = Model()

    def forward(self, x):
        return self.model(x)

    def example_arguments(self):
        return (torch.randn(2, 3, 224, 224),), None

    def input_names(self):
        return ("input",)

    def output_names(self):
        return ("output",)

    def dynamic_shapes(self):
        return ({0: "batch_size"},)


exportable_module = MyExportableModule()
onnx_program = exportable_module.to_onnx()
# The model can also be supplied directly to torch.onnx.export
onnx_program = torch.onnx.export(exportable_module)
dynamic_shapes() [исходный код]

Возвращает спецификации динамических форм входных данных модели.

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

Пример:

def dynamic_shapes(self):
    # Specify batch dimension as dynamic for input named 'x'
    return {"x": {0: "batch_size"}}


def dynamic_shapes(self):
    # Multiple dynamic dimensions
    return {
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "mask": {0: "batch_size"},
    }

Примечание

Реализация по умолчанию возвращает None, указывая, что все измерения являются статическими.

Возвращает:

Спецификация динамических форм, совместимая с torch.export.export. Возвращает None, если все измерения входных данных должны быть статическими. Формат может быть следующим:

  • Словарь, сопоставляющий имена входных данных спецификациям измерений
  • Кортеж/список спецификаций измерений, соответствующих входным данным
  • Любой формат, поддерживаемый параметром dynamic_shapes функции torch.export.export
Тип возвращаемого значения:

Any

abstract example_arguments() [исходный код]

Возвращает примеры аргументов для метода forward модели.

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

Пример:

def example_arguments(self):
    # For a model expecting a single tensor input
    return (torch.randn(1, 3, 224, 224),), None


def example_arguments(self):
    # For a model with multiple inputs and keyword arguments
    return (torch.randn(1, 3, 224, 224), torch.randn(1, 512)), {
        "temperature": 1.0
    }
Возвращает:
  • Кортеж позиционных аргументов, передаваемых методу forward
  • Словарь именованных аргументов (или None, если именованные аргументы не нужны)
Тип возвращаемого значения:

Кортеж, содержащий

input_names() [исходный код]

Возвращает имена входных тензоров модели.

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

Пример:

def input_names(self):
    return ["image", "mask"]


def input_names(self):
    # For a single input
    return ["input_tensor"]

Примечание

Реализация по умолчанию возвращает None, в результате чего имена генерируются автоматически.

Возвращает:

Последовательность строк, представляющих имена входных данных, или None, чтобы использовать имена по умолчанию. Количество имён должно совпадать с количеством позиционных аргументов метода forward.

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

Sequence[str] | None

output_names() [исходный код]

Возвращает имена выходных тензоров модели.

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

Пример:

def output_names(self):
    return ["logits", "probabilities"]


def output_names(self):
    # For a single output
    return ["prediction"]

Примечание

Реализация по умолчанию возвращает None, в результате чего имена генерируются автоматически.

Возвращает:

Последовательность строк, представляющих имена выходных данных, или None, чтобы использовать имена по умолчанию. Количество имён должно совпадать с количеством выходных данных метода forward. Для моделей, возвращающих несколько выходных данных, укажите имя для каждого из них.

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

Sequence[str] | None

to_onnx(**kwargs) [исходный код]

Экспортирует модуль в формат ONNX.

Этот метод представляет собой удобную обёртку над torch.onnx.export, которая автоматически использует заданные модулем примеры аргументов, динамические формы, а также имена входных и выходных данных. Дополнительные параметры экспорта можно указать с помощью именованных аргументов.

См. также: torch.onnx.export — полная документация по параметрам экспорта.

Параметры:

**kwargs (Any) –

Дополнительные именованные аргументы, передаваемые в torch.onnx.export. К часто используемым параметрам относятся:

  • opset_version (int): целевая версия opset ONNX
  • optimize (bool): применять ли оптимизации к экспортируемой модели
Возвращает:

Объект ONNXProgram, содержащий экспортированную модель и метаданные.

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

ONNXProgram

torch.onnx.is_in_onnx_export() [исходный код]

Возвращает, выполняется ли в данный момент экспорт в ONNX.

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

bool

class torch.onnx.OnnxExporterError

Ошибки, возникающие при экспорте в ONNX. Это базовый класс для всех ошибок экспортера.

class torch.onnx.InputObserver(value_if_missing=None)

Перехватывает метод forward для сбора входных и выходных данных. Эта информация используется для вывода динамических форм и аргументов экспорта.

Параметры:

value_if_missing (dict[str | int, Any] | None) – Если аргумент отсутствует, его значение по умолчанию будет взято из этого словаря. Это используется, когда после этапа предварительного заполнения аргумент исчезает (например, pixel_values), а другой добавляется (например, past_key_values). Значения используются только для вывода динамических форм и аргументов, а не для выполнения модели.

Примеры

>>> input_observer = InputObserver()
>>> with input_observer(model):
>>>     model(x1, y1)
>>>     model(x2, y2)
>>> ep = torch.export.export(  # or torch.onnx.export
>>>     model,
>>>     input_observer.infer_arguments(),
>>>     dynamic_shapes.input_observer.infer_dynamic_shapes(),
>>> )

Для LLM:

>>> input_observer = InputObserver()
>>> with input_observer(model):
>>>     model.generate(input_ids)
>>> ep = torch.export.export(  # or torch.onnx.export
>>>     model,
>>>     (),
>>>     kwargs=input_observer.infer_arguments(),
>>>     dynamic_shapes.input_observer.infer_dynamic_shapes(),
>>> )

В последнем примере рассматривается LLM, принимающая изображения и текст в качестве входных данных. Первый вызов метода forward, который мы пытаемся экспортировать, содержит pixel_values, но не содержит past_key_values. В следующих вызовах нет pixel_values, но есть past_key_values. Наблюдатель понимает, что нужны pixel_values и past_key_values, но они не обязательно должны быть указаны одновременно. Поскольку pixel_values появляется только в первом вызове, наблюдатель не может определить, как вывести пустой тензор для этого аргумента. Для этого и нужен аргумент value_if_missing. Следующий пример — не просто условный пример: он демонстрирует, как использовать его с transformers.

from transformers import pipeline

model_id = "tiny-random/gemma-3"
pipe = pipeline(
    "image-text-to-text",
    model=model_id,
    device="cpu",
    trust_remote_code=True,
    max_new_tokens=3,
    dtype=torch.float16,
)
messages = [
    {
        "role": "system",
        "content": [{"type": "text", "text": "You are a helpful assistant."}],
    },
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG",
            },
            {"type": "text", "text": "What animal is on the candy?"},
        ],
    },
]
observer = InputObserver(
    value_if_missing=dict(
        pixel_values=torch.empty((0, 3, 896, 896), dtype=torch.float16)
    )
)
with observer(pipe.model):
    pipe(text=messages, max_new_tokens=4)

Добавлено в версии 2.11.0.

check_discrepancies(onnx_program, atol=0.0001, rtol=0.1, progress_bar=False, initializer=<function _ort_session_initializer>, skip_none=True) [исходный код]

Вычисляет расхождения между сохранёнными входными и выходными данными и сохранённой моделью ONNX.

Параметры:
  • onnx_program (torch.onnx.ONNXProgram) – Экспортированная модель для проверки.
  • atol (float) – Абсолютный допуск; рекомендуемые значения: 1e-4 для float, 1e-2 для float16.
  • rtol (float) – Относительный допуск.
  • progress_bar (bool) – Отображает индикатор выполнения (требуется tqdm).
  • initializer (Callable[[str | bytes], ort.InferenceSession]) – Функция, вызываемая для инициализации сеанса инференса ONNX Runtime с указанной моделью. По умолчанию используется функция _ort_session_initializer.
  • skip_none (bool) – Не проверяет расхождения, если выходное значение равно None.
Возвращает:

Список словарей, готовый для передачи в dataframe.

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

list[dict[str, str | int | float | bool]]

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

infer_arguments(index_or_args_or_kwargs=None, flat=False, as_args_kwargs=False) [исходный код]

Выводит аргументы на основе собранных тензоров.

Параметры:
  • index_or_args_or_kwargs (tuple[Any] | dict[str, Any] | int | None) – Если значение не указано, метод выбирает один набор входных данных из доступных, обычно тот, который содержит наибольшее количество тензоров. Затем он заменяет значения None и отсутствующие тензоры пустыми тензорами. Если значение указано, это может быть целое число, позволяющее получить один из сохранённых наборов входных данных.
  • flat (bool) – Если True, возвращает плоский список тензоров; если False, возвращает кортеж или словарь, сохраняющий вложенную структуру. Плоская версия используется внутри метода. Она формирует единый список тензоров, который проще обрабатывать или изменять, чем вложенную структуру с теми же тензорами. Исходную структуру можно восстановить с помощью torch.utils._pytree.tree_unflatten(flat_list, self.aligned_spec). Этот механизм используется для замены значений None пустыми тензорами.
  • as_args_kwargs (bool) – Если True, метод всегда возвращает (args, kwargs); в противном случае возвращает либо кортеж (только args), либо словарь (только kwargs), либо вызывает исключение, если это невозможно.
Возвращает:

Выведенные аргументы, в которых каждый необязательный тензор заменён пустым тензором.

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

list[Tensor | None] | tuple[Tensor, …] | dict[str, Tensor] | tuple[list[Tensor] | tuple[Tensor, …], dict[str, Tensor]]

infer_dynamic_shapes(set_batch_dimension_for=None) [исходный код]

Выводит динамические формы. Обычно модели поддерживают измерение размера пакета, но это измерение имеет одинаковое значение для всех примеров входных данных. Вместо выполнения инференса на новых примерах можно использовать аргумент set_batch_dimension_for, чтобы указать, что первое измерение является динамическим для определённого набора входных данных, заданного их именами (str) или позициями (int).

Параметры:

set_batch_dimension_for (set[int | str] | bool | None) – Набор идентификаторов входных данных (по позиции, как int, или по имени, как str), для которых первое измерение следует считать динамическим измерением размера пакета. Если None, ни одно измерение не помечается явно как динамическое.

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

tuple[dict[int, Any] | None, …] | dict[str, dict[int, Any] | None]

num_obs() [исходный код]

Возвращает количество сохранённых наборов входных данных.

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

int

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/onnx_export.html

Spec-Zone.ru

Настройки Оффлайн Что нового Помощь О нас
Spec-Zone .ru
спецификации, руководства, описания, API