PyTorch 中如何将模型导出为 ONNX 格式？

推荐答案

在 PyTorch 中，可以使用 torch.onnx.export 函数将模型导出为 ONNX 格式。以下是一个简单的示例代码：

// www.javascriptcn.com code example
import torch
import torch.onnx

# 定义一个简单的模型
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = torch.nn.Linear(10, 1)

    def forward(self, x):
        return self.linear(x)

# 实例化模型
model = SimpleModel()

# 创建一个示例输入
dummy_input = torch.randn(1, 10)

# 导出模型为 ONNX 格式
torch.onnx.export(model,               # 要导出的模型
                  dummy_input,         # 模型输入
                  "model.onnx",        # 导出的文件名
                  export_params=True,  # 导出模型参数
                  opset_version=11,    # ONNX 操作集版本
                  do_constant_folding=True,  # 是否进行常量折叠优化
                  input_names=['input'],   # 输入名称
                  output_names=['output'], # 输出名称
                  dynamic_axes={'input': {0: 'batch_size'},  # 动态轴
                                'output': {0: 'batch_size'}})

本题详细解读

1. torch.onnx.export 函数参数详解

• model: 要导出的 PyTorch 模型。
• args: 模型的输入参数，通常是一个张量或张量的元组。
• f: 导出的 ONNX 文件的路径。
• export_params: 如果为 True，则导出模型的参数（权重和偏置）。如果为 False，则导出的模型将不包含参数。
• opset_version: 指定 ONNX 操作集的版本。不同的版本支持不同的操作。
• do_constant_folding: 是否进行常量折叠优化。常量折叠是一种优化技术，可以在导出时计算并折叠模型中的常量表达式。
• input_names: 输入张量的名称列表。
• output_names: 输出张量的名称列表。
• dynamic_axes: 指定动态轴的字典。动态轴允许输入和输出的某些维度在运行时变化。

2. ONNX 格式简介

ONNX（Open Neural Network Exchange）是一种开放的格式，用于表示深度学习模型。它允许模型在不同的深度学习框架之间进行转换和共享。ONNX 格式的模型可以在支持 ONNX 的推理引擎（如 ONNX Runtime、TensorRT 等）上运行。

3. 动态轴的使用

在 dynamic_axes 参数中，可以指定输入和输出的动态轴。例如，{'input': {0: 'batch_size'}} 表示输入张量的第 0 维是动态的，并且可以在运行时变化。这在处理不同批量大小的输入时非常有用。

4. 导出后的模型验证

导出 ONNX 模型后，可以使用 ONNX 的 onnx.checker 模块来验证模型的正确性：

import onnx

# 加载导出的 ONNX 模型
onnx_model = onnx.load("model.onnx")

# 验证模型
onnx.checker.check_model(onnx_model)

如果模型验证通过，说明导出的 ONNX 模型是有效的。