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

说实话，模型训练完了只是第一步。怎么把它从PyTorch搬到生产环境，这才是真正考验人的地方。我个人习惯用ONNX作为中间格式，再转TensorRT做推理加速。今天咱们就聊聊这个环节里最核心的一步——torch.onnx.export。

2.1 导出前的准备工作

先说说我踩过的坑。有一次我导出一个分类模型，怎么导都报错。折腾了半天，发现是模型里有个自定义的nn.Module，forward里用了torch.einsum。ONNX不支持这个操作。嗯，这里要注意：ONNX只支持标准算子。

所以在导出之前，我建议你先做两件事：

检查算子兼容性：用torch.onnx.export的check_trace参数做一次预检
固定输入形状：至少确定一个batch size，方便后续调试

小技巧：如果你模型里有torch.where、torch.topk这类操作，建议先换成ONNX支持的等价实现。我在项目中遇到过好几次，就是因为这些操作导致导出失败。

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

这个函数看起来参数很多，但真正常用的就那么几个。我按重要程度给你排个序：

参数名	作用	我的建议
`model`	要导出的PyTorch模型	记得先调成`eval()`模式
`args`	示例输入张量	形状要和实际推理时一致
`f`	输出ONNX文件路径	建议用`.onnx`后缀
`export_params`	是否导出模型参数	默认True，别改
`opset_version`	ONNX算子集版本	我一般用13或14，兼容性好
`do_constant_folding`	是否做常量折叠优化	建议开启，能减小模型体积
`input_names`	输入节点名称	给每个输入起个名字，方便后面用
`output_names`	输出节点名称	同理，方便TensorRT绑定
`dynamic_axes`	动态轴设置	这个最复杂，下面单独讲

来看一个最基础的导出示例：

import torch
import torchvision.models as models

# 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()

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

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=14,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output']
)

这段代码跑完，你就能得到一个resnet18.onnx文件。但说实话，这只是最基础的用法。真正生产环境里，你肯定需要处理动态轴。

2.3 动态轴与静态轴设置

什么叫动态轴？说白了就是允许输入的形状在推理时变化。比如你的模型要处理不同尺寸的图片，或者batch size不固定。

静态轴就简单了——形状完全固定。比如你只处理224x224的图片，batch size永远是1。那用静态轴就够了，性能还更好。

我个人习惯是：能静态就静态，需要动态才动态。因为动态轴会引入额外的性能开销，尤其是在TensorRT里。

2.3.1 静态轴设置

其实上面的例子就是静态轴。你传进去的dummy_input形状是(1, 3, 224, 224)，导出的ONNX就固定了这个形状。推理时输入必须完全匹配。

注意：静态轴虽然性能好，但灵活性差。我曾经给客户部署一个模型，他们用的图片尺寸是256x256，我导出的ONNX是224x224。结果跑不了，只能重新导出。所以一定要提前确认好输入规格。

2.3.2 动态轴设置

动态轴通过dynamic_axes参数来配置。这个参数是个字典，key是节点名称，value是轴索引和对应的动态范围。

来看一个例子：

# 设置动态轴
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

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

这里我做了三件事：

让input的batch size、高度、宽度都变成动态的
让output的batch size也跟着动态变化
给每个动态轴起了个名字，方便后面TensorRT里引用

你想想看，这样导出的ONNX就能接受任意batch size、任意尺寸的图片了。但代价是什么？性能会下降一些。因为TensorRT没法做针对固定形状的极致优化。

2.3.3 混合动态轴

还有一种情况：部分轴动态，部分轴静态。比如我只想让batch size动态，但图片尺寸固定：

dynamic_axes = {
    'input': {0: 'batch_size'},
    'output': {0: 'batch_size'}
}

这样导出的模型，batch size可以变，但图片必须是224x224。我在实际项目中经常用这种方案，既保留了灵活性，又不会损失太多性能。

2.4 验证导出的ONNX模型

导出完了别急着用，先验证一下。我一般用onnxruntime做推理测试：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
session = ort.InferenceSession("resnet18.onnx")

# 准备输入数据
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
outputs = session.run(None, {input_name: input_data})

print("ONNX推理成功，输出形状:", outputs[0].shape)

如果这一步能跑通，说明导出没问题。如果报错，多半是算子不兼容或者动态轴设置有问题。

  核心要点：

导出前检查算子兼容性

能静态就静态，需要动态才动态

动态轴用dynamic_axes字典配置

导出后用onnxruntime验证

2.5 知识体系总览

下面这张图总结了本章的核心逻辑，我画成了SVG，方便你理解整个流程：

这张图把整个导出流程串起来了。从准备工作开始，到核心参数配置，再到动态轴的选择，最后验证收尾。每一步都有坑，但只要你按这个流程走，基本不会出大问题。

好了，关于torch.onnx.export的核心参数和动态轴设置，就聊这么多。记住我强调的那句话：能静态就静态，需要动态才动态。下一章咱们会深入ONNX的算子兼容性问题，到时候再细聊。

可可精品资料库

📚 课程目录

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

2.1 导出前的准备工作

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

2.3 动态轴与静态轴设置

2.3.1 静态轴设置

2.3.2 动态轴设置

2.3.3 混合动态轴

2.4 验证导出的ONNX模型

2.5 知识体系总览