评价此页

★ ★ ★ ★ ★

torch.onnx#

创建日期：2025 年 6 月 10 日 | 最后更新日期：2025 年 9 月 10 日

概述#

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 引擎，以“预编译”（Ahead-of-Time, 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 节点中的可调用目标（例如，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 时有效。启用回退时，建议即使提供了 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 设置为一个具有以下结构的字典：

键（str）：输入或输出的名称。每个名称也必须在 input_names 或
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) – 已弃用。

返回:

如果 dynamo 为 True，则返回 torch.onnx.ONNXProgram，否则返回 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-block 时将其重置。

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

参数:

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

torch.onnx#

概述#

基于 torch.export 的 ONNX 导出器#

常见问题解答#

贡献/开发#

torch.onnx API#

函数#

类#

已弃用 API#

文档

教程

资源