评价此页

★ ★ ★ ★ ★

torch.onnx#

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

概述#

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

导出的模型可以被许多支持 ONNX 的运行时环境所使用，包括微软的 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 贡献指南，但您可能也会有兴趣阅读我们的开发维基。

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='.', 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 时

导出器尝试以下策略来获取用于转换为 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, ...]) – 示例位置输入参数。任何非 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) – 是否使用 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] | 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) – 是否对导出过程进行性能分析（profile）。此选项仅在 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)：输入或输出名称。每个名称还必须在 input_names 或
output_names 中提供。.
VALUE (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。

2.11 版本中的更改：fallback 选项已被移除。

torch.onnx.is_in_onnx_export()[源码]

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

返回类型:: 布尔值

类#

class torch.onnx.ONNXProgram(model, exported_program)

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

变量:

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

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

已弃用的 API#

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

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

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

当用户为自定义/第三方算子注册符号函数时，强烈建议通过 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)[源码]#

注销 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 – 类型和含义与 export() 的 model 参数相同。
mode (TrainingMode) – 类型和含义与 export() 的 training 参数相同。

torch.onnx#

概述#

基于 torch.export 的 ONNX 导出器#

常见问题解答#

贡献 / 开发#

torch.onnx API#

函数#

类#

已弃用的 API#

文档

教程

资源