ONNX简介 || 将PyTorch模型导出为ONNX || 扩展ONNX注册表

将PyTorch模型导出为ONNX

创建于：2023年10月04日 | 最后更新：2023年11月06日 | 最后验证：2024年11月05日

作者: Thiago Crepaldi

注意

截至 PyTorch 2.1，有两个版本的 ONNX 导出器。

• torch.onnx.dynamo_export 是基于 PyTorch 2.0 发布的 TorchDynamo 技术的最新（仍在测试阶段）导出器。

• torch.onnx.export 基于 TorchScript 后端，自 PyTorch 1.2.0 起可用。

在60 Minute Blitz中，我们有机会从高层次上学习PyTorch，并训练一个小型神经网络来分类图像。在本教程中，我们将扩展这一点，描述如何使用TorchDynamo和torch.onnx.dynamo_export 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 onnx
pip install onnxscript

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

一旦你的环境设置好了，让我们开始用PyTorch建模我们的图像分类器，就像我们在60 Minute Blitz中做的那样。

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)

/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:820: FutureWarning:

'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.

/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:820: FutureWarning:

'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.

/usr/local/lib/python3.10/dist-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.

/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/onnxfunction_dispatcher.py:503: FutureWarning:

'onnxscript.values.TracedOnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.

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

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电脑上，也可以直接从浏览器运行。 让我们尝试通过打开以下链接来使用网页版：https://netron.app/。

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

就是这样！我们成功地将我们的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模型，并将结果与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. Exporting a PyTorch model to ONNX

2. Extending the ONNX registry