torch.onnx

Функции

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 Node (например, 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) –

  Устарело: вместо этого задавайте dynamic_shapes при условии dynamo=True.

  По умолчанию формы всех входных и выходных тензоров экспортированной модели устанавливаются точно в соответствии с формами, указанными в 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: по умолчанию optimize теперь равен True.

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

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

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

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

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

bool

Классы

classtorch.onnx.ONNXProgram(model, exported_program)

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

Переменные:

• model – Модель ONNX в виде объекта модели ONNX IR.
• exported_program – Экспортированная программа, на основе которой создана модель ONNX.

classtorch.onnx.OnnxExporterError

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

Устаревшие API

torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version) [исходный код]

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

При регистрации пользователем символической функции для пользовательских операторов или операторов contrib настоятельно рекомендуется добавить для этого оператора вывод формы с помощью API setType; в противном случае в некоторых крайних случаях экспортированный граф может содержать неверный вывод формы. Пример использования setType — test_aten_embedding_2 в test_operators.py.

Пример использования см. в разделе «Пользовательские операторы» документации модуля.

Параметры:

• symbolic_name (str) – Имя пользовательского оператора в формате «<domain>::<op>».
• symbolic_fn (Callable) – Функция, принимающая граф ONNX и входные аргументы текущего оператора и возвращающая новые узлы операторов для добавления в граф.
• opset_version (int) – Версия набора операторов ONNX, в которой выполняется регистрация.

torch.onnx.unregister_custom_op_symbolic(symbolic_name, opset_version) [исходный код]

Отменяет регистрацию symbolic_name.

Пример использования см. в разделе «Пользовательские операторы» документации модуля.

Параметры:

• symbolic_name (str) – Имя пользовательского оператора в формате «<domain>::<op>».
• opset_version (int) – Версия набора операторов ONNX, регистрацию в которой следует отменить.

torch.onnx.select_model_mode_for_export(model, mode) [исходный код]

Менеджер контекста, который временно задаёт для model режим обучения mode и восстанавливает его при выходе из блока with.

Устарело начиная с версии 2.7: Задайте режим обучения перед экспортом модели.

Параметры:

• model – Тип и смысл совпадают с аргументом model функции export().
• mode (TrainingMode) – Тип и смысл совпадают с аргументом training функции export().