3、PyTorch模型导出为ONNX：torch.onnx.export详解，动态轴与静态轴设置，opset版本选择

好，咱们进入实战环节了。前面聊了ONNX的基础概念，现在要真刀真枪地干——把PyTorch模型转成ONNX。这一步，说白了就是给模型做一次「格式转换」，但里面的门道可不少。

我个人习惯把导出过程分成三步：准备模型、调用导出函数、验证结果。每一步都有坑，咱们一个一个来踩。

3.1 torch.onnx.export 核心参数

这个函数是PyTorch官方提供的导出接口。我刚开始用的时候，觉得参数太多记不住。后来发现，真正关键的其实就几个。

import torch
import torch.onnx

# 假设我们有一个训练好的模型
model = torchvision.models.resnet18(pretrained=True)
model.eval()

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出模型
torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 示例输入
    "resnet18.onnx",     # 输出文件名
    export_params=True,  # 是否导出参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],     # 输入节点名称
    output_names=['output'],   # 输出节点名称
    dynamic_axes=None          # 动态轴设置
)

嗯，这里要注意几个点。第一，模型一定要设成eval()模式。我在项目中遇到过有人忘了这步，结果导出的ONNX里还带着Dropout和BatchNorm的训练行为，推理时结果完全不对。

第二，dummy_input的形状要和实际输入一致。你想想看，如果模型输入是3通道224x224的图片，你给个1通道的，导出肯定报错。

3.2 动态轴与静态轴设置

这是最容易让人迷糊的地方。什么叫动态轴？说白了就是允许输入输出的某个维度在推理时变化。

举个例子。你训练了一个文本分类模型，输入是[batch_size, sequence_length]。如果batch_size固定为1，那就是静态轴。但如果你的服务需要同时处理多个请求，batch_size就得是动态的。

静态轴设置（默认）：所有维度都固定。导出时什么样，推理时就必须什么样。

# 静态轴：batch_size固定为1
torch.onnx.export(
    model,
    dummy_input,  # shape: [1, 3, 224, 224]
    "static.onnx",
    # 不设置dynamic_axes，所有轴都是静态的
)

动态轴设置：指定哪些维度可以变化。

# 动态轴：batch_size和sequence_length可变
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'sequence_length'},
    'output': {0: 'batch_size'}
}

torch.onnx.export(
    model,
    dummy_input,  # shape: [1, 3, 128]
    "dynamic.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes
)

我曾经踩过一个坑：给动态轴起的名字和模型内部变量名冲突了。ONNX导出时没报错，但推理时TensorRT死活跑不起来。排查了半天，才发现是命名问题。所以建议你给动态轴起名时，用batch_size、height、width这种语义清晰的词，别用dim0、dim1这种。

我的建议：除非你的应用场景明确需要动态batch，否则先用静态轴。动态轴会引入额外的性能开销，TensorRT优化时也会更复杂。

3.3 opset版本选择

opset，全称是Operator Set，就是ONNX支持的算子集合。每个版本会新增或废弃一些算子。

opset版本 发布时间 重要变化
9 2018年 支持动态形状，添加了ConstantOfShape等算子
11 2019年 添加了ScatterND、GatherElements等，修复了若干bug
13 2020年 支持Softmax的axis参数，添加了If、Loop等控制流算子
15 2021年 添加了Shape、Size等形状相关算子
17 2022年 支持LayerNormalization、GroupNormalization等

选哪个版本？我个人习惯用opset_version=11。为什么？因为这个版本足够稳定，TensorRT对它的支持也最成熟。我试过用opset 13导出，结果TensorRT报错说某个算子不支持，最后还得降级到11。

注意：opset版本不是越高越好。高版本虽然支持更多算子，但可能和推理引擎不兼容。建议先查一下目标推理引擎（比如TensorRT、ONNX Runtime）支持的opset范围。

如果你用了PyTorch的新特性，比如torch.where、torch.topk，可能需要高版本opset才能导出。这时候就得权衡一下：是改模型代码，还是升级opset版本。

3.4 导出后的验证

导出完别急着用，先验证一下。我一般做三件事：

  1. 检查ONNX模型结构：用onnx.load加载，打印计算图看看
  2. 运行ONNX Runtime推理：对比PyTorch和ONNX的输出是否一致
  3. 检查算子支持情况：用onnx.checker.check_model做完整性检查
import onnx
import onnxruntime as ort
import numpy as np

# 1. 加载并检查模型
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)

# 2. 运行ONNX Runtime推理
ort_session = ort.InferenceSession("resnet18.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)

# 3. 对比结果
torch_output = model(dummy_input).detach().numpy()
np.testing.assert_allclose(torch_output, ort_outputs[0], rtol=1e-03, atol=1e-05)
print("导出验证通过！")

嗯，这里有个小技巧。如果对比结果差异很大，别慌。先检查一下模型是否在eval()模式，再检查输入数据是否做了相同的预处理。我遇到过最离谱的一次，是忘了把输入从torch.float64转成torch.float32，结果ONNX推理结果完全不对。

3.5 常见问题与避坑指南

  • 算子不支持：PyTorch里有些算子ONNX没有对应实现。解决办法是改写模型，用ONNX支持的算子替代。比如torch.einsum可以用torch.matmul加torch.transpose组合实现。
  • 动态形状导出失败：如果模型里有torch.reshape或torch.view，动态轴可能会出问题。我建议在导出前，先用torch.jit.trace跟踪一下模型，看看有没有不支持动态形状的操作。
  • 导出文件太大：如果模型很大，ONNX文件可能几百MB。可以用onnx.compress_model压缩，或者用onnx.shape_inference.infer_shapes去掉冗余信息。

避坑指南：我曾经导出过一个带torch.nn.functional.interpolate的模型，opset 11下死活不成功。后来发现是align_corners参数的问题。解决办法是设置align_corners=False，或者升级到opset 13。所以，遇到问题先查官方文档，看看是不是opset版本的限制。

好了，这一章的内容就这些。总结一下：导出ONNX的核心是理解torch.onnx.export的参数含义，特别是动态轴和opset版本的选择。下一章咱们会讲如何用TensorRT优化导出的ONNX模型，让推理速度飞起来。

← 上一章 📖 返回目录 下一章 →