评价此页
★ ★ ★ ★ ★
beginner/onnx/export_simple_model_to_onnx_tutorial

注意

跳转至页面底部 下载完整示例代码。

ONNX 简介 || 将 PyTorch 模型导出为 ONNX || 扩展 ONNX 导出器的算子支持 || 将包含控制流的模型导出为 ONNX

将 PyTorch 模型导出为 ONNX#

创建日期：2023年10月4日 | 最后更新：2025年7月11日 | 最后验证：2024年11月5日

作者：Ti-Tai Wang, Justin Chu, Thiago Crepaldi。

注意

从 PyTorch 2.5 开始，有两种 ONNX 导出器选项可用。* torch.onnx.export(..., dynamo=True) 是推荐使用的导出器，它利用 torch.export 和 Torch FX 进行图捕获。* torch.onnx.export 是依赖于已弃用的 TorchScript 的传统方式，不再建议使用。

在60分钟闪电入门中，我们有机会从高层次了解了 PyTorch 并训练了一个小型神经网络来对图像进行分类。在本教程中，我们将对此进行扩展，介绍如何使用 torch.onnx.export(..., dynamo=True) ONNX 导出器将 PyTorch 定义的模型转换为 ONNX 格式。

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

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

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

  1. 安装所需的依赖项。

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

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

  4. 将 ONNX 模型保存到文件中。

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

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

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

1. 安装所需的依赖项#

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

pip install --upgrade onnx onnxscript

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

环境设置完成后，让我们开始用 PyTorch 构建图像分类器模型，过程与我们在60分钟闪电入门中所做的一样。

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


class ImageClassifierModel(nn.Module):
    def __init__(self):
        super().__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: torch.Tensor):
        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 = ImageClassifierModel()
# Create example inputs for exporting the model. The inputs should be a tuple of tensors.
example_inputs = (torch.randn(1, 1, 32, 32),)
onnx_program = torch.onnx.export(torch_model, example_inputs, dynamo=True)

/var/lib/workspace/beginner_source/onnx/export_simple_model_to_onnx_tutorial.py:93: UserWarning: Exporting a model while it is in training mode. Please ensure that this is intended, as it may lead to different behavior during inference. Calling model.eval() before export is recommended.
  onnx_program = torch.onnx.export(torch_model, example_inputs, dynamo=True)
[torch.onnx] Obtain model graph for `ImageClassifierModel([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `ImageClassifierModel([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decompositions...
/usr/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅

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

4. 将 ONNX 模型保存到文件中#

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

onnx_program.save("image_classifier_model.onnx")

你可以使用以下代码将 ONNX 文件加载回内存，并检查其格式是否正确：

import onnx

onnx_model = onnx.load("image_classifier_model.onnx")
onnx.checker.check_model(onnx_model)

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

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

../../_images/netron_web_ui.png

Netron 打开后，我们可以将 image_classifier_model.onnx 文件拖放到浏览器中，或者在点击 Open model（打开模型）按钮后选择该文件。 ../../_images/image_classifier_onnx_model_on_netron_web_ui.png

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

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

最后一步是使用 ONNX Runtime 执行 ONNX 模型，但在操作之前，我们需要安装 ONNX Runtime。

pip install onnxruntime

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

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

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

import onnxruntime

onnx_inputs = [tensor.numpy(force=True) for tensor in example_inputs]
print(f"Input length: {len(onnx_inputs)}")
print(f"Sample input: {onnx_inputs}")

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

onnxruntime_input = {input_arg.name: input_value for input_arg, input_value in zip(ort_session.get_inputs(), onnx_inputs)}

# ONNX Runtime returns a list of outputs
onnxruntime_outputs = ort_session.run(None, onnxruntime_input)[0]

Input length: 1
Sample input: [array([[[[ 0.81520087,  1.4223824 , -2.3707998 , ..., -1.2884982 ,
           1.0608014 ,  0.25451723],
         [ 0.3446546 ,  2.1509576 , -1.0615014 , ..., -0.5936997 ,
          -0.01130752,  0.5120359 ],
         [ 1.2144618 ,  1.6991181 , -1.3451372 , ...,  0.53668165,
          -0.15685211, -0.28054774],
         ...,
         [-0.641897  ,  0.19047514,  0.45556557, ...,  0.84515625,
           1.1896348 , -0.39858592],
         [-1.3440754 , -1.0720246 , -0.5851963 , ..., -0.86487633,
          -0.6251062 ,  1.1385156 ],
         [ 0.52466315, -1.0250294 , -2.4209156 , ..., -0.59231746,
           0.31369153,  0.25885025]]]], dtype=float32)]

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

确定导出模型是否良好的最佳方法是通过与我们的事实来源（即 PyTorch）进行数值评估。

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

torch_outputs = torch_model(*example_inputs)

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: [[ 0.06747029  0.02339114  0.06199165 -0.03377082 -0.0108053  -0.05977552
  -0.01948961 -0.03325256 -0.06009166  0.04576406]]

结论#

以上就是全部内容！我们已成功将 PyTorch 模型导出为 ONNX 格式，将模型保存到磁盘，使用 Netron 进行了查看，使用 ONNX Runtime 进行了执行，并最终将其数值结果与 PyTorch 的结果进行了比较。

延伸阅读#

下面的列表引用了从基本示例到高级场景的教程，不一定按列出的顺序。您可以随时跳转到您感兴趣的特定主题，或者坐下来，享受学习 ONNX 导出器所有知识的乐趣。

1. 将 PyTorch 模型导出到 ONNX
2. 扩展 ONNX 导出器算子支持
3. 将带控制流的模型导出到 ONNX

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

在此页面上
PyTorch 库