基于 TorchDynamo 的 ONNX 导出器#
创建于:2025 年 6 月 10 日 | 最后更新于:2025 年 6 月 20 日
概述#
ONNX 导出器利用 TorchDynamo 引擎挂钩 Python 的帧求值 API,并将 Python 字节码动态地重写为 FX 图。生成的 FX 图随后经过优化,最终转换为 ONNX 图。
这种方法的优势在于,FX 图是通过字节码分析捕获的,它保留了模型的动态特性,而不是使用传统的静态跟踪技术。
此外,在导出过程中,与 TorchScript 导出的模型相比,内存使用量大大降低。有关更多信息,请参阅内存使用情况文档。
依赖项#
ONNX 导出器依赖额外的 Python 包
可以通过 pip 安装它们。
pip install --upgrade onnx onnxscript
然后可以使用 onnxruntime 在各种处理器上执行模型。
一个简单的示例#
下面以一个简单的多层感知器(MLP)为例,演示导出器 API 的使用。
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 IR 规范的 onnx.ModelProto
对象。然后可以使用 torch.onnx.ONNXProgram.save()
API 将 ONNX 模型序列化为 Protobuf 文件。
onnx_program.save("mlp.onnx")
使用与 TorchScript 导出的模型相同的模型进行比较#
TorchScript 导出的模型和基于 TorchDynamo 的导出器之间的最大区别在于,后者要求 dynamic_shapes 与输入具有相同的树状结构,而前者则要求 dynamic_shapes 是单个扁平化的字典。
torch.onnx.export(model,(tensor_input, dict_input, list_input), "mlp.onnx", dynamic_axes={"tensor_input":{0: "batch_size"}, "tensor_x": {0: "batch_size"}, "list_input_index_0": {0: "batch_size"}}, input_names=input_names, output_names=output_names)
元数据#
在 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 op、参数和关键字参数。
示例:
%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)
这些元数据字段存储在每个 ONNX 节点的 metadata_props 属性中,并可以使用 Netron 或以编程方式进行检查。
整个 ONNX 图具有以下 metadata_props
pkg.torch.export.ExportedProgram.graph_signature
此属性包含原始 PyTorch ExportedProgram 中 graph_signature 的字符串表示形式。图签名描述了模型输入和输出的结构以及它们如何映射到 ONNX 图。输入定义为
InputSpec
对象,其中包含输入的种类(例如,参数的InputKind.PARAMETER
,用户定义输入的InputKind.USER_INPUT
)、参数名称、目标(可以是模型中的特定节点)以及输入是否持久。输出定义为OutputSpec
对象,指定输出的种类(例如,OutputKind.USER_OUTPUT
)和参数名称。要详细了解图签名,请参阅 torch.export 以获取更多信息。
pkg.torch.export.ExportedProgram.range_constraints
此属性包含原始 PyTorch ExportedProgram 中存在的任何范围约束的字符串表示形式。范围约束指定模型中符号形状或值的有效范围,这对于使用动态形状或符号维度的模型可能很重要。
示例:
s0: VR[2, int_oo]
,表示输入张量的尺寸必须至少为 2。要详细了解范围约束,请参阅 torch.export 以获取更多信息。
ONNX 图中的每个输入值可能具有以下元数据属性
pkg.torch.export.graph_signature.InputSpec.kind
输入种类,由 PyTorch 的 InputKind 枚举定义。
示例值
“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
输入种类,由 PyTorch 的 OutputKind 枚举定义。
示例值
“USER_OUTPUT”:用户可见的输出。
“LOSS_OUTPUT”:损失值输出。
“BUFFER_MUTATION”:指示缓冲区已被修改。
“GRADIENT_TO_PARAMETER”:参数的梯度输出。
“GRADIENT_TO_USER_INPUT”:用户输入的梯度输出。
“USER_INPUT_MUTATION”:指示用户输入已被修改。
“TOKEN”:令牌输出。
每个初始值、输入、输出都有以下元数据
pkg.torch.onnx.original_node_name
PyTorch FX 图中生成此值的原始节点名称,以防该值被重命名。这有助于将初始化器追溯到原始模型中的来源。
示例:
fc1.weight
API 参考#
- torch.onnx.export(model, args=(), f=None, *, kwargs=None, export_params=True, verbose=None, input_names=None, output_names=None, opset_version=None, dynamic_axes=None, keep_initializers_as_inputs=False, dynamo=False, 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, 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
和一套更现代的转换逻辑。这是导出模型的推荐方式。当
dynamo=True
时导出器尝试以下策略来获取用于转换为 ONNX 的 ExportedProgram。
如果模型已经是 ExportedProgram,则将按原样使用。
使用
torch.export.export()
并设置strict=False
。使用
torch.export.export()
并设置strict=True
。使用
draft_export
,它会删除数据相关操作中的一些健全性保证,以允许导出继续进行。如果导出器遇到任何不健全的数据相关操作,您将收到警告。使用
torch.jit.trace()
来跟踪模型,然后转换为 ExportedProgram。这是最不健全的策略,但可能有助于将 TorchScript 模型转换为 ONNX。
- 参数
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”。
export_params (bool) – 如果为 false,则不会导出参数(权重)。
verbose (bool | None) – 是否启用详细日志记录。
input_names (Sequence[str] | None) – 要分配给图中输入节点的名称,按顺序。
output_names (Sequence[str] | None) – 要分配给图中输出节点的名称,按顺序。
opset_version (int | None) – 要针对的 默认(ai.onnx)opset 的版本。必须 >= 7。
dynamic_axes (Mapping[str, Mapping[int, str]] | Mapping[str, Sequence[int]] | None) –
默认情况下,导出的模型的所有输入和输出张量的形状都将精确匹配
args
中给出的形状。要将张量的轴指定为动态(即,仅在运行时已知),请将dynamic_axes
设置为具有以下架构的字典- KEY (str): 输入或输出名称。每个名称也必须在
input_names
或 output_names 中提供
.
- KEY (str): 输入或输出名称。每个名称也必须在
- 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 ...
keep_initializers_as_inputs (bool) –
如果为 True,则导出的图中的所有初始化器(通常对应于模型权重)也将作为输入添加到图中。如果为 False,则初始化器不会作为输入添加到图中,只有用户输入会作为输入添加。
如果打算在运行时提供模型权重,请将其设置为 True。如果权重是静态的,请将其设置为 False,以便后端/运行时进行更好的优化(例如,常量折叠)。
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。
training (_C_onnx.TrainingMode) – 已弃用选项。而是,在导出之前设置模型的训练模式。
operator_export_type (_C_onnx.OperatorExportTypes) – 已弃用选项。仅支持 ONNX。
do_constant_folding (bool) – 已弃用选项。
custom_opsets (Mapping[str, int] | None) –
已弃用。一个字典
KEY (str): opset 域名称
VALUE (int): opset 版本
如果
model
引用了自定义 opset 但未在此字典中提及,则 opset 版本设置为 1。只能通过此参数指示自定义 opset 域名称和版本。export_modules_as_functions (bool | Collection[type[torch.nn.Module]]) –
已弃用选项。
标志,用于将所有
nn.Module
前向调用导出为 ONNX 中的本地函数。或者是一个集合,用于指定要导出为本地函数的特定类型的模块。此功能需要opset_version
>= 15,否则导出将失败。这是因为opset_version
< 15 意味着 IR 版本 < 8,因此不支持本地函数。模块变量将导出为函数属性。函数属性有两种类别。1. 标注属性:具有通过 PEP 526 样式进行的类型注解的类变量将作为属性导出。由于标注属性不是由 PyTorch JIT 跟踪创建的,因此它们不会在 ONNX 本地函数的子图中被使用,但消费者可能会使用它们来确定是否将函数替换为特定的融合内核。
2. 推断属性:被模块内的运算符使用的变量。属性名称将带有前缀“inferred::”。这是为了与从 Python 模块注解中检索的预定义属性区分开来。推断属性在 ONNX 本地函数的子图中使用。
False
(默认):将nn.Module
前向调用导出为细粒度节点。True
:将所有nn.Module
前向调用导出为本地函数节点。- nn.Module 类型集合:将
nn.Module
前向调用导出为本地函数节点, 仅当
nn.Module
的类型在此集合中时。
- nn.Module 类型集合:将
autograd_inlining (bool) – 已弃用。用于控制是否内联 autograd 函数的标志。有关详细信息,请参阅 pytorch/pytorch#74765。
- 返回
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。
- class torch.onnx.ONNXProgram(model, exported_program)#
表示可使用 torch 张量调用的 ONNX 程序的类。
- 变量
model – ONNX 模型,作为 ONNX IR 模型对象。
exported_program – 生成 ONNX 模型的导出程序。
- apply_weights(state_dict)[source]#
将指定 state_dict 中的权重应用于 ONNX 模型。
使用此方法替换 FakeTensors 或其他权重。
- 参数
state_dict (dict[str, torch.Tensor]) – 包含要应用于 ONNX 模型的权重的 state_dict。
- compute_values(value_names, args=(), kwargs=None)[source]#
计算 ONNX 模型中指定名称的值。
此方法用于计算 ONNX 模型中指定名称的值。这些值以将名称映射到张量的字典形式返回。
- 参数
value_names (Sequence[str]) – 要计算的值的名称。
- 返回
将名称映射到张量的字典。
- 返回类型
Sequence[torch.Tensor]
- initialize_inference_session(initializer=<function _ort_session_initializer>)[source]#
初始化 ONNX Runtime 推理会话。
- property model_proto: ModelProto#
返回 ONNX
ModelProto
对象。
- save(destination, *, include_initializers=True, keep_initializers_as_inputs=False, external_data=None)[source]#
将 ONNX 模型保存到指定的目标。
当
external_data
为True
或模型大于 2GB 时,权重将作为外部数据保存在单独的文件中。初始化器(模型权重)序列化行为
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 | os.PathLike) – 保存 ONNX 模型的路径。
include_initializers (bool) – 是否将初始化器包含在保存的模型中。
keep_initializers_as_inputs (bool) – 是否将初始化器保留为已保存模型中的输入。如果为 True,则将初始化器添加为模型的输入,这意味着它们可以被覆盖,方法是将其作为模型输入提供。
external_data (bool | None) – 是否将权重保存为单独文件中的外部数据。
- 引发
TypeError – 如果
external_data
为True
且destination
不是文件路径。
- class torch.onnx.OnnxExporterError#
ONNX 导出器引发的错误。这是所有导出器错误的基类。
- torch.onnx.enable_fake_mode()[source]#
在上下文期间启用伪模式。
在内部,它实例化一个
torch._subclasses.fake_tensor.FakeTensorMode
上下文管理器,该管理器将用户输入和模型参数转换为torch._subclasses.fake_tensor.FakeTensor
。伪张量 (
torch._subclasses.fake_tensor.FakeTensor
) 是一个torch.Tensor
,它能够在不实际进行通过meta
设备分配的张量进行计算的情况下运行 PyTorch 代码。由于设备上没有实际分配数据,因此此 API 允许在没有实际执行所需的内存占用的情况下初始化和导出大型模型。强烈建议在导出无法放入内存的大型模型时,在伪模式下初始化模型。
注意
此函数不支持 torch.onnx.export(…, dynamo=True, optimize=True)。请在模型导出后在函数外部调用 ONNXProgram.optimize()。
示例
# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX) >>> import torch >>> class MyModel(torch.nn.Module): # Model with a parameter ... def __init__(self) -> None: ... super().__init__() ... self.weight = torch.nn.Parameter(torch.tensor(42.0)) ... def forward(self, x): ... return self.weight + x >>> with torch.onnx.enable_fake_mode(): ... # When initialized in fake mode, the model's parameters are fake tensors ... # They do not take up memory so we can initialize large models ... my_nn_module = MyModel() ... arg1 = torch.randn(2, 2, 2) >>> onnx_program = torch.onnx.export(my_nn_module, (arg1,), dynamo=True, optimize=False) >>> # Saving model WITHOUT initializers (only the architecture) >>> onnx_program.save( ... "my_model_without_initializers.onnx", ... include_initializers=False, ... keep_initializers_as_inputs=True, ... ) >>> # Saving model WITH initializers after applying concrete weights >>> onnx_program.apply_weights({"weight": torch.tensor(42.0)}) >>> onnx_program.save("my_model_with_initializers.onnx")
警告
此 API 是实验性的,并且*不*向后兼容。
已弃用#
以下类和函数已弃用,将被移除。
- torch.onnx.dynamo_export(model, /, *model_args, export_options=None, **model_kwargs)[source]#
将 torch.nn.Module 导出为 ONNX 图。
版本 2.7 已弃用: 请改用
torch.onnx.export(..., dynamo=True)
。- 参数
model (torch.nn.Module | Callable | torch.export.ExportedProgram) – 要导出为 ONNX 的 PyTorch 模型。
model_args –
model
的位置输入。model_kwargs –
model
的关键字输入。export_options (ExportOptions | None) – 影响导出到 ONNX 的选项。
- 返回
导出的 ONNX 模型的内存中表示。
- 返回类型