torch.onnx

概述 ¶

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

导出的模型可以被任何支持 ONNX 的许多运行时消费，包括微软的 ONNX 运行时。

您可以使用以下列出的两种 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 语句或循环；

• 不处理 training 和 eval 模式之间的细微差别；

• 真正无法处理动态输入

为了支持静态追踪的限制，导出器还支持 TorchScript 脚本（通过 torch.jit.script() ），这增加了对数据相关控制流的支持，例如。然而，TorchScript 本身是 Python 语言的一个子集，因此并非所有 Python 的功能都得到支持，例如就地操作。

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

贡献/开发 ¶

ONNX 导出器是一个社区项目，我们欢迎贡献。我们遵循 PyTorch 的贡献指南，但你也许还愿意阅读我们的开发维基。