如何将PyTorch模型（.pth）转换为ONNX格式？

Question

**问题：如何将PyTorch模型（.pth）转换为ONNX格式？** 在深度学习模型部署过程中，常常需要将训练好的PyTorch模型（通常保存为`.pth`或`.pt`文件）转换为ONNX（Open Neural Network Exchange）格式，以便在不同框架或推理引擎中进行跨平台使用。然而，在实际操作中，开发者常遇到诸如模型结构复杂、输入维度不匹配、算子不支持等问题，导致导出失败。本文将详细介绍如何将PyTorch模型成功转换为ONNX格式，涵盖模型加载、模型定义、dummy input构建、torch.onnx.export的参数设置以及常见错误排查等内容，帮助你顺利完成模型转换。

Answer 1

如何将PyTorch模型（.pth）转换为ONNX格式？

在深度学习模型部署过程中，常常需要将训练好的PyTorch模型（通常保存为.pth或.pt文件）转换为ONNX（Open Neural Network Exchange）格式，以便在不同框架或推理引擎中进行跨平台使用。然而，在实际操作中，开发者常遇到诸如模型结构复杂、输入维度不匹配、算子不支持等问题，导致导出失败。

1. 环境准备与依赖安装

在开始转换之前，请确保已安装以下必要的库：

• torch：用于加载和运行PyTorch模型
• onnx：用于验证ONNX模型的正确性

可以通过以下命令安装所需库：

pip install torch onnx

2. 加载PyTorch模型

首先，我们需要加载训练好的PyTorch模型。通常有两种方式保存模型：

1. 仅保存模型参数（推荐）：torch.save(model.state_dict(), 'model.pth')
2. 保存整个模型对象：torch.save(model, 'model.pth')

根据保存方式的不同，加载方法也略有差异：

# 方法一：加载模型参数
model = MyModel()
model.load_state_dict(torch.load('model.pth'))

# 方法二：加载整个模型
model = torch.load('model.pth')

注意：如果你使用了自定义模型类，必须提前定义好该类，并且保持结构一致。

3. 构建Dummy Input

在导出ONNX模型时，需要一个示例输入来追踪模型计算图。这个输入应与模型期望的输入形状一致。

# 假设模型输入是 batch_size x channels x height x width
dummy_input = torch.randn(1, 3, 224, 224)

对于多输入模型，可以传入元组：

dummy_input1 = torch.randn(1, 3, 224, 224)
dummy_input2 = torch.randn(1, 10)
dummy_inputs = (dummy_input1, dummy_input2)

4. 使用torch.onnx.export导出模型

PyTorch提供了torch.onnx.export()函数用于导出ONNX模型。其基本调用如下：

torch.onnx.export(
    model,
    dummy_input,
    "output_model.onnx",
    export_params=True,  # 存储训练参数
    opset_version=13,  # ONNX算子集版本
    do_constant_folding=True,  # 优化常量
    input_names=['input'],  # 输入节点名称
    output_names=['output'],  # 输出节点名称
    dynamic_axes={
        'input': {0: 'batch_size'},  # 动态维度
        'output': {0: 'batch_size'}
    }
)

参数名	说明
export_params	是否存储训练参数，默认为True
opset_version	ONNX算子集版本，建议选择最新稳定版如13或17
dynamic_axes	指定动态维度，便于支持变长输入

5. 验证ONNX模型

导出完成后，建议使用onnx.checker验证模型的合法性：

import onnx

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

还可以使用onnxruntime进行前向推理测试：

import onnxruntime as ort

ort_session = ort.InferenceSession("output_model.onnx")
outputs = ort_session.run(
    None,
    {'input': dummy_input.numpy()}
)

6. 常见问题及解决方案

在转换过程中可能会遇到一些常见错误，以下是典型问题及其解决办法：

1. 错误：Input shape mismatch
   原因：dummy_input的shape与模型期望的输入不符。
   解决：检查模型输入定义，确保dummy_input的维度匹配。
2. 错误：Unsupported operator
   原因：某些PyTorch算子尚未支持导出到ONNX。
   解决：升级PyTorch版本；尝试使用torchscript先转换为ScriptModule。
3. 错误：Can't export Python ops
   原因：模型中包含自定义Python函数。
   解决：将Python操作替换为可导出的PyTorch等价操作。

7. 高级技巧与注意事项

为了提升模型兼容性和性能，可以考虑以下进阶做法：

• 使用TorchScript预处理模型
  将模型转换为TorchScript格式有助于更稳定的导出过程：

script_model = torch.jit.script(model)
torch.jit.save(script_model, "script_model.pt")

• 动态输入支持
  在dynamic_axes中合理设置动态维度，使模型支持任意batch size或其他维度变化。
• ONNX优化工具
  使用onnxoptimizer对导出后的ONNX模型进行优化，减少冗余计算。

8. 总结

通过上述步骤，我们可以将PyTorch模型成功转换为ONNX格式，并在多种推理引擎中使用。掌握这一流程不仅有助于模型部署，也为后续的模型压缩、量化、加速等操作打下基础。