TorchDynamo-based ONNX Exporter¶
警告
用于 TorchDynamo 的 ONNX 导出器是一项快速发展的 Beta 技术。
概述¶
ONNX 导出器利用 TorchDynamo 引擎连接到 Python 的帧评估 API，并动态地将其字节码重写为 FX Graph。然后对生成的 FX Graph 进行润色，最后将其转换为 ONNX 图。
这种方法的主要优点是 FX graph 是使用字节码分析捕获的，它保留了模型的动态特性，而不是使用传统的静态追踪技术。
此外，在导出过程中，与启用 TorchScript 的导出器相比，内存使用量显著减少。有关更多信息，请参阅 文档。
依赖项¶
ONNX 导出器依赖于额外的 Python 包
它们可以通过 pip 安装
pip install --upgrade onnx onnxscript
onnxruntime 然后可以用于在各种处理器上执行模型。
一个简单的例子¶
请参阅以下示例，演示了导出器 API 在一个简单的多层感知器 (MLP) 示例中的应用
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)
def forward(self, tensor_x: torch.Tensor):
tensor_x = self.fc0(tensor_x)
tensor_x = torch.sigmoid(tensor_x)
tensor_x = self.fc1(tensor_x)
tensor_x = torch.sigmoid(tensor_x)
tensor_x = self.fc2(tensor_x)
tensor_x = torch.sigmoid(tensor_x)
output = self.fc3(tensor_x)
return output
model = MLPModel()
tensor_x = torch.rand((97, 8), dtype=torch.float32)
onnx_program = torch.onnx.export(model, (tensor_x,), dynamo=True)
如上面的代码所示，您只需要提供 torch.onnx.export() 模型实例及其输入即可。然后，导出器将返回一个 torch.onnx.ONNXProgram 实例，其中包含导出的 ONNX 图以及额外信息。
可以调用 onnx_program.optimize() 来优化具有常量折叠和消除冗余运算符的 ONNX 图。优化是就地完成的。
onnx_program.optimize()
通过 onnx_program.model_proto 提供的内存模型是一个 onnx.ModelProto 对象，符合 ONNX IR 规范。然后可以使用 torch.onnx.ONNXProgram.save() API 将 ONNX 模型序列化为 Protobuf 文件。
onnx_program.save("mlp.onnx")
存在两个函数可以基于 TorchDynamo 引擎将模型导出到 ONNX。它们在生成 torch.export.ExportedProgram 的方式上略有不同。torch.onnx.dynamo_export() 是在 PyTorch 2.1 中引入的，而 torch.onnx.export() 在 PyTorch 2.5 中得到扩展，以便轻松地从 TorchScript 切换到 TorchDynamo。要调用前一个函数，可以将前一个示例的最后一行替换为以下行。
注意
torch.onnx.dynamo_export() 将在未来版本中弃用。请改用带有参数 dynamo=True 的 torch.onnx.export()。
onnx_program = torch.onnx.dynamo_export(model, tensor_x)
API 参考¶
- torch.onnx.dynamo_export(model, /, *model_args, export_options=None, **model_kwargs)[source][source]¶
将 torch.nn.Module 导出到 ONNX 图。
- 参数
model (torch.nn.Module | Callable | torch.export.ExportedProgram) – 要导出到 ONNX 的 PyTorch 模型。
model_args –
model的位置输入。model_kwargs –
model的关键字输入。export_options (ExportOptions | None) – 影响导出到 ONNX 的选项。
- 返回
导出的 ONNX 模型的内存表示。
- 返回类型
示例 1 - 最简单的导出
class MyModel(torch.nn.Module): def __init__(self) -> None: super().__init__() self.linear = torch.nn.Linear(2, 2) def forward(self, x, bias=None): out = self.linear(x) out = out + bias return out model = MyModel() kwargs = {"bias": 3.0} args = (torch.randn(2, 2, 2),) onnx_program = torch.onnx.dynamo_export(model, *args, **kwargs).save( "my_simple_model.onnx" )
示例 2 - 使用动态形状导出
# The previous model can be exported with dynamic shapes export_options = torch.onnx.ExportOptions(dynamic_shapes=True) onnx_program = torch.onnx.dynamo_export( model, *args, **kwargs, export_options=export_options ) onnx_program.save("my_dynamic_model.onnx")
- class torch.onnx.ONNXProgram(model, exported_program)¶
一个表示 ONNX 程序的类，该程序可以使用 torch 张量调用。
- apply_weights(state_dict)[source][source]¶
将指定状态字典中的权重应用于 ONNX 模型。
使用此方法替换 FakeTensors 或其他权重。
- 参数
state_dict (dict[str, torch.Tensor]) – 包含要应用于 ONNX 模型的权重的状态字典。
- initialize_inference_session(initializer=<function _ort_session_initializer>)[source][source]¶
初始化 ONNX Runtime 推理会话。
- property model_proto: ModelProto¶
返回 ONNX
ModelProto对象。
- save(destination, *, include_initializers=True, keep_initializers_as_inputs=False, external_data=None)[source][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) – 是否将权重作为外部数据保存在单独的文件中。
- Raises
TypeError – 如果
external_data为True且destination不是文件路径。
- class torch.onnx.ExportOptions(*, dynamic_shapes=None, fake_context=None, onnx_registry=None, diagnostic_options=None)¶
影响 TorchDynamo ONNX 导出器的选项。
- 变量
dynamic_shapes (bool | None) – 输入/输出张量的形状信息提示。当为
None时，导出器确定最兼容的设置。当为True时，所有输入形状都被认为是动态的。当为False时，所有输入形状都被认为是静态的。diagnostic_options (DiagnosticOptions) – 导出器的诊断选项。
fake_context (ONNXFakeContext | None) – 用于符号追踪的 fake 上下文。
onnx_registry (OnnxRegistry | None) – 用于将 ATen 运算符注册到 ONNX 函数的 ONNX 注册表。
- torch.onnx.enable_fake_mode()[source]¶
在上下文期间启用 fake 模式。
在内部，它实例化一个
torch._subclasses.fake_tensor.FakeTensorMode上下文管理器，该管理器将用户输入和模型参数转换为torch._subclasses.fake_tensor.FakeTensor。torch._subclasses.fake_tensor.FakeTensor是一个torch.Tensor，它能够在不实际通过在meta设备上分配的张量进行计算的情况下运行 PyTorch 代码。由于设备上没有分配实际数据，因此此 API 允许初始化和导出大型模型，而无需执行它所需的实际内存占用。强烈建议在导出过大而无法放入内存的模型时，在 fake 模式下初始化模型。
示例
# 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) >>> # 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 是实验性的，不 向后兼容。
- class torch.onnx.ONNXRuntimeOptions(*, session_options=None, execution_providers=None, execution_provider_options=None)¶
影响通过 ONNX Runtime 执行 ONNX 模型的选项。
- class torch.onnx.OnnxExporterError¶
ONNX 导出器引发的错误。这是所有导出器错误的基本类。
- class torch.onnx.OnnxRegistry¶
ONNX 函数的注册表。
注册表维护从限定名称到固定 opset 版本下的符号函数的映射。它支持注册自定义 onnx-script 函数，并供调度程序调度对相应函数的调用。
- get_op_functions(namespace, op_name, overload=None)[source][source]¶
返回给定运算符的 ONNXFunctions 列表：torch.ops.<namespace>.<op_name>.<overload>。
列表按注册时间排序。自定义运算符应位于列表的后半部分。
- is_registered_op(namespace, op_name, overload=None)[source][source]¶
返回给定的 op 是否已注册：torch.ops.<namespace>.<op_name>.<overload>。
- register_op(function, namespace, op_name, overload=None, is_complex=False)[source][source]¶
注册自定义操作符：torch.ops.<namespace>.<op_name>.<overload>。
- 参数
- Raises
ValueError – 如果名称格式不是 ‘namespace::op’。
