4. PyTorch模型导出ONNX：torch.onnx.export详解、动态轴与静态轴设置

模型训练完了，总不能一直窝在PyTorch里跑吧？尤其是做YOLO推理，咱们得把它弄到ONNX里，再转成TensorRT，才能真正跑起来。这一步，说白了就是给模型换个“壳”，让它能在不同的引擎里跑。

我个人习惯，导出ONNX是整套流程里最容易踩坑的一步。你想想看，PyTorch里模型跑得好好的，一导出就报错，或者导出来精度不对，这种事我遇到过太多次了。今天咱们就把torch.onnx.export这个函数掰开揉碎了讲清楚。

4.1 torch.onnx.export 核心参数

这个函数，参数不少，但真正关键的其实就几个。我列个表，你一看就明白。

参数名	类型	作用	我的建议
model	torch.nn.Module	要导出的PyTorch模型	记得先调成eval()模式
args	tuple	模型的输入张量	形状要和实际推理时一致
f	str	输出ONNX文件路径	用.onnx后缀
export_params	bool	是否导出模型权重	默认True，别改
opset_version	int	ONNX算子集版本	YOLO建议用11或更高
do_constant_folding	bool	是否做常量折叠优化	推理用True，训练用False
input_names	list[str]	输入节点名称	随便起，但要有意义
output_names	list[str]	输出节点名称	TensorRT里会用到
dynamic_axes	dict	定义动态轴	这是重点，下面细说

4.2 静态轴 vs 动态轴

这个区别，我刚开始也迷糊过。说白了，静态轴就是固定死的，动态轴就是可以变的。

拿YOLOv8举例，输入一般是[1, 3, 640, 640]。这里的1是batch size。如果你只做单张图片推理，那batch size固定为1就行，这就是静态轴。但如果你要一次处理多张图片，比如[4, 3, 640, 640]，那batch size就得是动态的。

为什么会这样？因为ONNX在导出时，如果指定了静态轴，它会针对那个固定形状做很多优化。动态轴则保留灵活性，但优化空间会小一些。

4.3 dynamic_axes 配置详解

这个参数是个字典，结构稍微有点绕。我给你拆开看。

# 假设输入输出名称
input_names = ['images']
output_names = ['output0']

# 动态轴配置
dynamic_axes = {
    'images': {0: 'batch_size'},      # 输入的batch维度动态
    'output0': {0: 'batch_size'}      # 输出的batch维度动态
}

# 导出调用
torch.onnx.export(
    model,
    dummy_input,
    'yolov8n.onnx',
    input_names=input_names,
    output_names=output_names,
    dynamic_axes=dynamic_axes,
    opset_version=11
)

你看，dynamic_axes的key是节点名称，value是另一个字典。那个字典的key是轴索引（从0开始），value是给这个轴起的名字。名字随便起，但最好有意义。

如果你有多个输出，每个输出的动态轴都要单独配。我曾经漏配了一个输出头的动态轴，结果推理时batch size一变，那个头就报错。

4.4 完整导出示例（YOLOv8）

咱们直接上代码。这是我常用的YOLOv8导出脚本，你拿去改改就能用。

import torch
import torch.onnx

def export_onnx(model_path, output_path, img_size=640, batch_size=1, dynamic=True):
    # 加载模型
    model = torch.load(model_path, map_location='cpu')['model'].float()
    model.eval()
    
    # 创建虚拟输入
    dummy_input = torch.randn(batch_size, 3, img_size, img_size)
    
    # 定义输入输出名称
    input_names = ['images']
    output_names = ['output0']
    
    # 配置动态轴
    if dynamic:
        dynamic_axes = {
            'images': {0: 'batch_size'},
            'output0': {0: 'batch_size'}
        }
    else:
        dynamic_axes = None
    
    # 导出
    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        opset_version=11,
        do_constant_folding=True,
        export_params=True
    )
    
    print(f'导出成功: {output_path}')

# 使用示例
export_onnx('yolov8n.pt', 'yolov8n.onnx', dynamic=True)

4.5 导出后的验证

导出来不代表就对了。我每次导出完，都会用ONNX Runtime跑一遍，和PyTorch的结果对比一下。

import onnx
import onnxruntime as ort
import numpy as np

# 检查ONNX模型结构
onnx_model = onnx.load('yolov8n.onnx')
onnx.checker.check_model(onnx_model)
print('ONNX模型结构检查通过')

# 用ONNX Runtime推理
ort_session = ort.InferenceSession('yolov8n.onnx')
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

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

# 推理
ort_outputs = ort_session.run([output_name], {input_name: dummy_input})

# 对比结果（这里简化了，实际要比较数值差异）
print(f'ONNX Runtime输出形状: {ort_outputs[0].shape}')

这一步不能省。我见过有人导出的ONNX精度掉了0.5个点，就是因为某个算子不兼容。早点发现，早点排查。

4.6 常见问题与排查

导出过程中，你可能会遇到这些问题。我列个表，方便你对照。

问题	原因	解决方案
算子不支持	PyTorch算子没有对应的ONNX实现	升级opset版本，或替换算子
动态轴报错	动态轴配置与模型实际结构不匹配	检查所有输入输出的轴索引
精度下降	常量折叠或算子精度问题	关闭do_constant_folding试试
导出卡死	模型太大或内存不足	减小输入尺寸，或分批导出
形状不匹配	动态轴导致输出形状变化	在模型中固定某些维度

嗯，到这里，ONNX导出的核心内容就差不多了。你只要把torch.onnx.export的参数搞明白，动态轴和静态轴的区别弄清楚，导出这一步基本不会出大问题。下一节，咱们就要开始碰TensorRT了，那才是真正考验耐心的地方。