3、PyTorch模型导出ONNX：torch.onnx.export核心参数详解、动态轴与静态轴设置

好，咱们进入正题。模型训练完了，总不能一直活在PyTorch的温室里吧？得把它放到ONNX这个“中间格式”里，才能去嵌入式设备上跑。这一步，说白了就是给模型办个“护照”，让它在不同框架间自由通行。

我个人习惯把导出ONNX看作一次“模型快照”。你不仅要拍下它的长相（网络结构），还得记下它的脾气（输入输出的形状、数据类型）。torch.onnx.export这个函数，就是干这个活的。

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

先看一个最基础的调用，我带你拆开揉碎了讲。

import torch
import torchvision.models as models

# 1. 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()  # 切换到推理模式，这点很重要

# 2. 构造一个 dummy input（虚拟输入）
#    形状是 [batch_size, channels, height, width]
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出
torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 模型的输入张量
    "resnet18.onnx",     # 输出的ONNX文件名
    export_params=True,  # 是否导出训练好的参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 是否执行常量折叠优化
    input_names=['input'],     # 输入节点的名字
    output_names=['output'],   # 输出节点的名字
    dynamic_axes=None          # 动态轴设置，后面细讲
)

这里每个参数都有它的脾气，我挑几个重点说说。

3.1.1 export_params

这个参数默认就是True。你想想看，如果只导出网络结构不带权重，那跟画了个空壳子有什么区别？我在项目中遇到过有人忘了设这个，结果导出的模型推理结果全是乱码。嗯，排查了半天才发现是权重没带出来。

3.1.2 opset_version

这个参数决定了ONNX支持哪些算子。版本越高，支持的算子越新，但兼容性可能差一些。我个人习惯用opset_version=11，这个版本比较成熟，大部分嵌入式设备的推理引擎都支持。如果你用了很新的PyTorch算子，比如一些变形金刚里的操作，可能需要升到13或更高。

3.1.3 do_constant_folding

这个参数我建议你打开。它会把一些在推理时不会变化的计算（比如1+1这种）提前算好，直接塞进模型里。说白了就是帮你做一次“预编译”，减少推理时的计算量。对性能有好处，而且不会改变模型精度。

3.1.4 input_names 和 output_names

给输入输出节点起个名字。这看起来是小事，但实际部署时特别重要。比如你在C++的推理引擎里加载模型，就得靠这个名字来绑定数据。我习惯用'input'和'output'，简单明了。如果你有多个输入，可以写成['input1', 'input2']。

3.2 动态轴与静态轴设置

这是最容易踩坑的地方，也是面试官最爱问的点。什么是动态轴？说白了，就是允许模型在推理时接受不同大小的输入。

举个例子，你训练时用的图片都是224x224，但实际部署时，用户上传的图片可能是300x400。如果你的模型是静态的，那就只能处理224x224，大了小了都得报错。动态轴就是解决这个问题的。

3.2.1 静态轴（Static Axes）

默认情况下，所有轴都是静态的。也就是说，你导出时用的dummy_input是什么形状，推理时就必须是什么形状。比如上面那个例子，dummy_input是[1, 3, 224, 224]，那推理时batch size只能是1，图片尺寸只能是224x224。

静态轴的好处是性能好。推理引擎可以针对固定尺寸做极致优化，比如内存预分配、算子融合等。如果你的应用场景输入尺寸固定，那就用静态轴，省心又高效。

3.2.2 动态轴（Dynamic Axes）

动态轴允许你在推理时改变某些维度的大小。最常见的场景就是batch size可变，或者图片尺寸可变。

看代码：

# 设置动态轴
dynamic_axes = {
    'input': {
        0: 'batch_size',   # 第0维（batch）是动态的
        2: 'height',       # 第2维（高度）是动态的
        3: 'width'         # 第3维（宽度）是动态的
    },
    'output': {
        0: 'batch_size'    # 输出的batch也是动态的
    }
}

torch.onnx.export(
    model,
    dummy_input,
    "resnet18_dynamic.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes,
    opset_version=11
)

这里dynamic_axes是一个字典，key是输入输出节点的名字，value又是一个字典，里面指定了哪些维度是动态的，以及给这个维度起个名字（比如'batch_size'）。

3.2.3 什么时候用动态轴？

我总结了几种场景：

• batch size可变： 这是最常见的。比如你的服务端要同时处理多张图片，但每次请求的图片数量不一样。把batch size设成动态，就不用为每个batch size都导出一个模型了。
• 图片尺寸可变： 比如做目标检测，输入图片的原始尺寸可能不一样。但要注意，模型里如果有Resize操作，那还好说；如果有全连接层，那就得小心了。
• 序列长度可变： 处理文本或语音时，输入序列的长度经常是变化的。这时候必须用动态轴。

3.2.4 避坑指南

另外，动态轴的名字（比如'batch_size'）只是给人看的，对推理引擎没有实际意义。你起什么名字都行，只要保证输入和输出中对应的维度名字一致即可。

3.3 验证导出的ONNX模型

导出完了，别急着部署。先验证一下模型是否正确。ONNX官方提供了一个工具叫onnxruntime，可以用来跑一下推理，看看结果跟PyTorch是否一致。

import onnx
import onnxruntime as ort
import numpy as np

# 加载ONNX模型
onnx_model = onnx.load("resnet18.onnx")
# 检查模型结构是否合法
onnx.checker.check_model(onnx_model)

# 创建推理会话
ort_session = ort.InferenceSession("resnet18.onnx")

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
outputs = ort_session.run(
    None,
    {'input': input_data}
)

# 对比PyTorch结果
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(input_data))

# 检查误差
np.testing.assert_allclose(
    torch_outputs.numpy(), 
    outputs[0], 
    rtol=1e-03, 
    atol=1e-05
)
print("模型导出验证通过！")

这里rtol和atol是相对误差和绝对误差的容忍度。一般来说，1e-3和1e-5就够用了。如果误差太大，说明导出过程中可能出了问题，比如某些算子不支持，被替换成了精度较低的实现。

3.4 小结

嗯，这一章的内容就这些。总结一下：

• torch.onnx.export的核心参数要记牢，特别是opset_version和do_constant_folding。
• 动态轴和静态轴的选择，取决于你的应用场景。能静态就别动态，这是性能的保证。
• 导出后一定要用onnxruntime验证，确保精度无损。

下一章，我们会讲如何对导出的ONNX模型进行优化和剪枝，让它更适合在嵌入式设备上跑。到时候见。