3、模型导出实战：从PyTorch导出ONNX模型，torch.onnx.export参数详解

好，咱们直接进入正题。

前面两章我们把ONNX的底子打好了，这一章就来点真格的——把PyTorch模型转成ONNX。说白了，就是让PyTorch里的模型，能跑到别的框架或者硬件上去。

我个人习惯，每次做模型导出前，都会先问自己三个问题：

• 这个模型是干啥的？分类、检测还是分割？
• 导出后要跑在哪个后端？TensorRT、OpenVINO还是手机端？
• 输入输出的形状和类型，是不是固定的？

想清楚这些，再动手写代码，能少踩很多坑。

3.1 最简单的导出：一个分类模型

咱们先来个最简单的例子。假设你训练了一个ResNet-18分类模型，想把它导出成ONNX。

代码其实就几行：

import torch
import torchvision.models as models

# 1. 加载模型，设为eval模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=11
)

嗯，这里要注意：模型一定要设成eval模式。我在项目中遇到过有人忘了这步，结果导出的ONNX里包含了Dropout和BatchNorm的训练行为，推理结果完全不对。你想想看，这坑踩得多冤。

3.2 torch.onnx.export 核心参数详解

上面那个例子只是开胃菜。真正干活的时候，你得把参数吃透。我挑几个最常用的，一个一个说。

参数名	作用	我的建议
model	要导出的PyTorch模型	必须设成eval模式
args	虚拟输入（tuple或tensor）	形状和类型要和真实推理时一致
f	输出文件名	建议用.onnx后缀
input_names	输入节点的名字	用有意义的名称，比如"input"、"image"
output_names	输出节点的名字	同理，比如"output"、"logits"
opset_version	ONNX算子集版本	我一般用11或13，兼容性好
dynamic_axes	动态轴设置	处理可变batch size时必用
export_params	是否导出模型参数	默认True，别改
verbose	是否打印导出日志	调试时设为True，能看细节

3.3 动态轴：处理可变batch size

大多数场景下，你的输入batch size是固定的。但有些时候，比如在线服务，你希望一次能处理不同数量的图片。这时候就需要动态轴了。

torch.onnx.export(
    model,
    dummy_input,
    "resnet18_dynamic.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=11
)

这里我解释一下：dynamic_axes是个字典，key是节点名，value是维度映射。比如{0: "batch_size"}，意思就是第0维（batch维）是动态的，名字叫"batch_size"。

我曾经在做一个视频分析项目时，需要每帧都跑模型，但视频分辨率不固定。当时没设动态轴，结果导出的ONNX只能处理固定尺寸，后来改了半天才搞定。嗯，从那以后，只要输入尺寸可能变化，我必加dynamic_axes。

3.4 opset_version：选对版本很重要

ONNX的算子集是不断演进的。不同版本支持的算子不一样。选错了，导出可能失败，或者推理时出问题。

我建议，除非你有特殊需求，否则先用11。等跑通了，再尝试更高的版本。别一上来就搞个17，结果发现TensorRT不支持，那就尴尬了。

3.5 常见坑与避坑指南

导出ONNX看着简单，但实际项目中坑不少。我把自己踩过的坑列出来，你遇到了直接对照着查。

3.6 实战：导出一个带后处理的模型

很多时候，模型不只是输出一个tensor，还要做一些后处理，比如softmax、argmax。这些操作能不能也导进ONNX？

答案是：可以，但不建议。

我个人习惯，只导出模型的核心推理部分，后处理放在外面做。原因有二：

1. 后处理算子（比如argmax）在不同后端上实现可能不同，容易出兼容性问题。
2. 后处理逻辑经常变，如果写在ONNX里，每次改都要重新导出，太麻烦。

但如果你非要导，也不是不行。比如这样：

class ModelWithSoftmax(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, x):
        x = self.model(x)
        x = self.softmax(x)
        return x

model_with_softmax = ModelWithSoftmax(model)
model_with_softmax.eval()

torch.onnx.export(
    model_with_softmax,
    dummy_input,
    "resnet18_softmax.onnx",
    input_names=["input"],
    output_names=["probabilities"],
    opset_version=11
)

你看，把softmax包在模型里，导出的ONNX就直接输出概率值了。但说实话，我很少这么干。除非是给非技术人员用的工具，他们不想再写后处理代码。

3.7 验证导出的ONNX模型

导出完了，别急着用。先验证一下，确保模型没问题。

import onnx
import onnxruntime as ort
import numpy as np

# 1. 检查模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("模型结构检查通过")

# 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. 对比PyTorch和ONNX的输出
with torch.no_grad():
    torch_output = model(dummy_input)

np.testing.assert_allclose(
    torch_output.numpy(), ort_outputs[0], rtol=1e-03, atol=1e-05
)
print("输出一致性检查通过")

这一步非常关键。我见过太多人导出完就扔一边，结果部署时才发现精度对不上。你想想看，要是上线了才发现，那得多被动。

3.8 小结

这一章我们干了三件事：

• 学会了最简单的模型导出
• 吃透了torch.onnx.export的核心参数
• 掌握了验证导出的方法

下一章，我们会深入ONNX的内部结构，看看它到底长什么样。到时候，你就能亲手修改ONNX模型了。