快捷方式

torch.onnx

概览

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

导出的模型可以被许多支持 ONNX 的运行时所消费,包括 Microsoft 的 ONNX Runtime

您可以使用如下所列的两种 ONNX 导出器 API。两者都可以通过函数 torch.onnx.export() 调用。下一个示例展示了如何导出简单模型。

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
)

接下来的章节介绍了导出器的两个版本。

基于 TorchDynamo 的 ONNX 导出器

基于 TorchDynamo 的 ONNX 导出器是适用于 PyTorch 2.1 及更新版本的最新的(Beta)导出器

利用 TorchDynamo 引擎挂接到 Python 的帧评估 API,并将其字节码动态重写为 FX 图。生成的 FX 图经过精修后,最终被转换为 ONNX 图。

这种方法的主要优势在于,FX 图是使用字节码分析捕获的,这保留了模型的动态特性,而不是使用传统的静态跟踪技术。

了解更多关于基于 TorchDynamo 的 ONNX 导出器

基于 TorchScript 的 ONNX 导出器

基于 TorchScript 的 ONNX 导出器自 PyTorch 1.2.0 版本起可用

利用 TorchScript 跟踪模型(通过 torch.jit.trace())并捕获静态计算图。

因此,生成的图存在一些限制

  • 它不记录任何控制流,例如 if 语句或循环;

  • 不处理训练模式和评估模式之间的细微差别;

  • 不能真正处理动态输入

为了尝试支持静态跟踪的限制,导出器还支持 TorchScript 脚本化(通过 torch.jit.script()),例如,这增加了对数据依赖的控制流的支持。然而,TorchScript 本身是 Python 语言的一个子集,因此并非所有 Python 特性都受支持,例如原地操作。

了解更多关于基于 TorchScript 的 ONNX 导出器

贡献 / 开发

ONNX 导出器是一个社区项目,我们欢迎贡献。我们遵循PyTorch 的贡献指南,但也建议您阅读我们的开发维基

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源