torch.onnx#

创建于: 2025年6月10日 | 最后更新于: 2025年8月29日

概述#

Open Neural Network eXchange (ONNX) 是一种表示机器学习模型的开放标准格式。torch.onnx 模块捕获原生 PyTorch torch.nn.Module 模型中的计算图，并将其转换为 ONNX 图。

导出的模型可以被许多支持 ONNX 的运行时使用，包括 Microsoft 的 ONNX Runtime。

下一个示例展示了如何导出简单模型。

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 128, 5)

    def forward(self, x):
        return torch.relu(self.conv1(x))

input_tensor = torch.rand((1, 1, 128, 128), dtype=torch.float32)

model = MyModel()

torch.onnx.export(
    model,                  # model to export
    (input_tensor,),        # inputs of the model,
    "my_model.onnx",        # filename of the ONNX model
    input_names=["input"],  # Rename inputs for the ONNX model
    dynamo=True             # True or False to select the exporter to use
)

基于 torch.export 的 ONNX 导出器#

基于 torch.export 的 ONNX 导出器是 PyTorch 2.6 及更新版本的最新导出器

使用 torch.export 引擎以“提前编译”（AOT）方式生成一个仅表示函数张量计算的追踪图。生成的追踪图（1）在功能性的 ATen 算子集中产生规范化的算子（以及任何用户指定的自定义算子），（2）消除了所有 Python 控制流和数据结构（某些例外情况除外），并且（3）记录了显示这种规范化和控制流消除对未来输入是健全的形状约束集，然后才最终将其翻译成 ONNX 图。

了解更多关于基于 torch.export 的 ONNX 导出器

常见问题#

问：我已经导出了我的 LLM 模型，但它的输入大小似乎是固定的？

追踪器会记录示例输入的形状。如果模型应该接受动态形状的输入，请在调用 torch.onnx.export() 时设置 dynamic_shapes。

问：如何导出包含循环的模型？

请参阅 torch.cond。

贡献/开发#

ONNX 导出器是一个社区项目，我们欢迎贡献。我们遵循 PyTorch 的贡献指南，但您可能也对阅读我们的开发 Wiki 感兴趣。

torch.onnx 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='.', fallback=False, 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)[source]

将模型导出为 ONNX 格式。

设置 dynamo=True 将启用新的 ONNX 导出逻辑，该逻辑基于 torch.export.ExportedProgram 和更现代的翻译逻辑集。这是将模型导出为 ONNX 的推荐方法，也是默认方法。

当 dynamo=True 时

导出器尝试以下策略来获取用于转换为 ONNX 的 ExportedProgram。

如果模型已经是 ExportedProgram，则按原样使用。
使用 torch.export.export() 并设置 strict=False。
使用 torch.export.export() 并设置 strict=True。

参数

model (torch.nn.Module | torch.export.ExportedProgram | torch.jit.ScriptModule | torch.jit.ScriptFunction) – 要导出的模型。
args (tuple[Any, ...]) – 示例位置输入。任何非张量参数都将被硬编码到导出的模型中；任何张量参数都将成为导出模型的输入，顺序与元组中的顺序相同。
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) – 是否使用 torch.export ExportedProgram 而不是 TorchScript 来导出模型。
external_data (bool) – 是否将模型权重保存为外部数据文件。这对于权重超过 ONNX 文件大小限制（2GB）的模型是必需的。当设置为 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 | Sequence[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 时有效。
fallback (bool) – 如果 dynamo 导出器失败，是否回退到 TorchScript 导出器。此选项仅在 dynamo 为 True 时有效。启用 fallback 时，建议即使提供了 dynamic_shapes，也设置 dynamic_axes。
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 且 fallback 未启用时，优先指定 dynamic_shapes。

默认情况下，导出的模型的所有输入和输出张量的形状都将精确匹配 args 中给出的形状。要将张量的轴指定为动态（即仅在运行时才知道），请将 dynamic_axes 设置为具有以下架构的字典：

KEY (str)：输入或输出名称。每个名称还必须在 input_names 或
output_names 中提供。.
VALUE (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。

torch.onnx.is_in_onnx_export()[source]

返回当前是否处于 ONNX 导出过程中。

返回类型: 布尔值

类#

class torch.onnx.ONNXProgram(model, exported_program)

一个表示可与 torch 张量调用的 ONNX 程序的类。

变量

model – ONNX 模型，作为 ONNX IR 模型对象。
exported_program – 生成 ONNX 模型的导出程序。

class torch.onnx.OnnxExporterError: ONNX 导出器引发的错误。这是所有导出器错误的基类。

已弃用的 API#

版本 2.6 已弃用：这些函数已被弃用，将在未来版本中删除。

torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)[source]#

为自定义运算符注册符号函数。

当用户为自定义/贡献运算符注册符号时，强烈建议通过 setType API 为该运算符添加形状推断，否则导出的图在某些极端情况下可能具有不正确的形状推断。setType 的一个示例是 test_operators.py 中的 test_aten_embedding_2。

有关示例用法，请参阅模块文档中的“自定义运算符”。

参数

symbolic_name (str) – 自定义运算符的名称，格式为“<domain>::<op>”。
symbolic_fn (Callable) – 一个函数，它接收 ONNX 图和当前运算符的输入参数，并返回要添加到图中的新运算符节点。
opset_version (int) – 要注册的 ONNX 算子集版本。

torch.onnx.unregister_custom_op_symbolic(symbolic_name, opset_version)[source]#

注销 symbolic_name。

有关示例用法，请参阅模块文档中的“自定义运算符”。

参数

symbolic_name (str) – 自定义运算符的名称，格式为“<domain>::<op>”。
opset_version (int) – 要注销的 ONNX 算子集版本。

torch.onnx.select_model_mode_for_export(model, mode)[source]#

一个上下文管理器，用于临时将 model 的训练模式设置为 mode，并在退出 with 块时将其重置。

版本 2.7 已弃用：请在导出模型之前设置训练模式。

参数

model – 与 export() 的 model 参数类型和含义相同。
mode (TrainingMode) – 与 export() 的 training 参数类型和含义相同。

torch.onnx#

概述#

基于 torch.export 的 ONNX 导出器#

常见问题#

贡献/开发#

torch.onnx API#

函数#

类#

已弃用的 API#

文档

教程

资源