注意

单击此处下载完整的示例代码

ONNX 入门 || 将 PyTorch 模型导出到 ONNX || 扩展 ONNX 注册表

将 PyTorch 模型导出到 ONNX¶

注意

从 PyTorch 2.1 开始，ONNX 导出器有两个版本。

torch.onnx.dynamo_export 是基于 TorchDynamo 技术的最新（仍处于测试阶段）导出器，该技术与 PyTorch 2.0 一起发布
torch.onnx.export 基于 TorchScript 后端，自 PyTorch 1.2.0 起可用

在 60 分钟速成中，我们有机会在高层次上了解 PyTorch 并训练一个小型的图像分类神经网络。在本教程中，我们将扩展这一点，描述如何使用 TorchDynamo 和 torch.onnx.dynamo_export ONNX 导出器将 PyTorch 中定义的模型转换为 ONNX 格式。

虽然 PyTorch 非常适合迭代模型开发，但模型可以使用不同的格式部署到生产环境，包括 ONNX (开放神经网络交换)！

ONNX 是一种灵活的开放标准格式，用于表示机器学习模型，它标准化了机器学习的表示形式，使它们能够跨从大型云超级计算机到资源受限的边缘设备（如您的网络浏览器和手机）等各种硬件平台和运行时环境执行。

在本教程中，我们将学习如何

安装所需的依赖项。
编写一个简单的图像分类器模型。
将模型导出到 ONNX 格式。
将 ONNX 模型保存在文件中。
使用 Netron 可视化 ONNX 模型图。
使用 ONNX Runtime 执行 ONNX 模型
比较 PyTorch 结果和 ONNX Runtime 结果。

1. 安装所需的依赖项¶

由于 ONNX 导出器使用 onnx 和 onnxscript 将 PyTorch 算子转换为 ONNX 算子，因此我们需要安装它们。

pip install onnx
pip install onnxscript

2. 编写一个简单的图像分类器模型¶

环境设置完成后，让我们用 PyTorch 开始对我们的图像分类器进行建模，就像我们在 60 分钟速成教程中做的那样。

import torch
import torch.nn as nn
import torch.nn.functional as F


class MyModel(nn.Module):

    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

3. 将模型导出为 ONNX 格式¶

现在我们已经定义了模型，我们需要实例化它并创建一个随机的 32x32 输入。接下来，我们可以将模型导出为 ONNX 格式。

torch_model = MyModel()
torch_input = torch.randn(1, 1, 32, 32)
onnx_program = torch.onnx.dynamo_export(torch_model, torch_input)

/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/onnx/_internal/_exporter_legacy.py:116: UserWarning:

torch.onnx.dynamo_export only implements opset version 18 for now. If you need to use a different opset version, please register them with register_custom_op.

正如我们所看到的，我们不需要对模型进行任何代码更改。生成的 ONNX 模型存储在 torch.onnx.ONNXProgram 中，作为一个二进制 protobuf 文件。

4. 将 ONNX 模型保存在文件中¶

虽然在内存中加载导出的模型在许多应用程序中非常有用，但我们可以使用以下代码将其保存到磁盘

onnx_program.save("my_image_classifier.onnx")

您可以使用以下代码将 ONNX 文件重新加载到内存中，并检查其是否格式正确

import onnx
onnx_model = onnx.load("my_image_classifier.onnx")
onnx.checker.check_model(onnx_model)

5. 使用 Netron 可视化 ONNX 模型图¶

现在我们已经将模型保存到文件中，我们可以使用 Netron 对其进行可视化。Netron 可以安装在 macos、Linux 或 Windows 计算机上，也可以直接从浏览器运行。让我们尝试通过打开以下链接来使用 Web 版本：https://netron.app/。

Netron 打开后，我们可以将 my_image_classifier.onnx 文件拖放到浏览器中，或者在点击“打开模型”按钮后选择它。

../../_images/image_clossifier_onnx_modelon_netron_web_ui.png

就这样！我们已经成功地将 PyTorch 模型导出为 ONNX 格式，并使用 Netron 对其进行了可视化。

6. 使用 ONNX Runtime 执行 ONNX 模型¶

最后一步是使用 ONNX Runtime 执行 ONNX 模型，但在我们这样做之前，让我们安装 ONNX Runtime。

pip install onnxruntime

ONNX 标准不支持 PyTorch 支持的所有数据结构和类型，因此我们需要在将 PyTorch 输入传递给 ONNX Runtime 之前将其适应为 ONNX 格式。在我们的示例中，输入恰好相同，但在更复杂的模型中，它可能比原始 PyTorch 模型具有更多输入。

ONNX Runtime 需要一个额外的步骤，该步骤涉及将所有 PyTorch 张量转换为 Numpy（在 CPU 上），并将它们包装在一个字典中，该字典的键是一个字符串，包含输入名称作为键，Numpy 张量作为值。

现在我们可以创建一个ONNX Runtime 推理会话，使用处理后的输入执行 ONNX 模型，并获取输出。在本教程中，ONNX Runtime 在 CPU 上执行，但它也可以在 GPU 上执行。

import onnxruntime

onnx_input = onnx_program.adapt_torch_inputs_to_onnx(torch_input)
print(f"Input length: {len(onnx_input)}")
print(f"Sample input: {onnx_input}")

ort_session = onnxruntime.InferenceSession("./my_image_classifier.onnx", providers=['CPUExecutionProvider'])

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

onnxruntime_input = {k.name: to_numpy(v) for k, v in zip(ort_session.get_inputs(), onnx_input)}

onnxruntime_outputs = ort_session.run(None, onnxruntime_input)

Input length: 1
Sample input: (tensor([[[[-1.0416,  1.1125, -0.3602,  ..., -0.0189, -1.2205,  0.3472],
          [ 0.6651,  1.1037, -0.3674,  ..., -1.4724,  0.2539, -0.0788],
          [-0.1239, -0.6458, -0.7785,  ..., -0.2674,  0.3019, -0.5682],
          ...,
          [-0.0300, -0.4833, -0.3928,  ..., -1.2406,  0.8488, -0.5473],
          [-0.8185, -0.1276,  0.3475,  ..., -1.0702, -1.6922, -0.6048],
          [ 0.8268, -0.0248, -0.3354,  ..., -0.9178, -0.3240,  0.7485]]]]),)

7. 将 PyTorch 结果与 ONNX Runtime 的结果进行比较¶

确定导出模型是否良好，最好的方法是通过与 PyTorch 进行数值评估，而 PyTorch 是我们的真实来源。

为此，我们需要使用相同的输入执行 PyTorch 模型，并将结果与 ONNX Runtime 的结果进行比较。在比较结果之前，我们需要将 PyTorch 的输出转换为与 ONNX 格式匹配的格式。

torch_outputs = torch_model(torch_input)
torch_outputs = onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)

assert len(torch_outputs) == len(onnxruntime_outputs)
for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
    torch.testing.assert_close(torch_output, torch.tensor(onnxruntime_output))

print("PyTorch and ONNX Runtime output matched!")
print(f"Output length: {len(onnxruntime_outputs)}")
print(f"Sample output: {onnxruntime_outputs}")

PyTorch and ONNX Runtime output matched!
Output length: 1
Sample output: [array([[ 0.14531787, -0.05903321, -0.00652155,  0.09054166,  0.01458297,
        -0.08046442, -0.12109031, -0.03938238, -0.01814789, -0.01363543]],
      dtype=float32)]

结论¶

就是这样！我们已经成功地将 PyTorch 模型导出为 ONNX 格式，将模型保存到磁盘，使用 Netron 对其进行了可视化，使用 ONNX Runtime 对其进行了执行，最后将它的数值结果与 PyTorch 的进行了比较。

进一步阅读¶

以下列表引用了从基本示例到高级场景的教程，不一定是按照它们列出的顺序。您可以直接跳到您感兴趣的特定主题，也可以安心地浏览所有主题，以了解有关 ONNX 导出器的所有知识。

1. 将 PyTorch 模型导出到 ONNX

2. 扩展 ONNX 注册表

脚本的总运行时间：（0 分钟 1.533 秒）

Sphinx-Gallery 生成的图库