📚 课程目录第4章 ONNX模型导出与标准化

第4章：ONNX模型导出与标准化

好，咱们进入实战环节的第一个关键步骤——ONNX导出。很多同学在PyTorch里跑得飞起的模型，一到TensorRT就各种报错，十有八九是ONNX这步没处理好。我见过太多人花了两周调精度，最后发现是导出时一个参数设错了。

说白了，ONNX就是PyTorch和TensorRT之间的翻译官。翻译得好不好，直接决定你后面的部署能不能顺利。

4.1 torch.onnx.export核心参数

先看最常用的导出方式。我习惯把参数分成三组：必填项、动态形状相关、算子控制相关。

import torch

# 假设你有一个训练好的模型
model = YourDetectionModel()
model.eval()

# 构造一个 dummy input，形状要和实际推理时一致
dummy_input = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 示例输入
    "model.onnx",             # 输出路径
    export_params=True,       # 导出训练好的权重
    opset_version=11,         # 算子集版本
    do_constant_folding=True, # 常量折叠优化
    input_names=["input"],    # 输入张量名称
    output_names=["output"],  # 输出张量名称
    dynamic_axes={            # 动态轴设置
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

关键参数速查表

参数	推荐值	说明
opset_version	11 或 13	工业检测建议用11，兼容性最好
do_constant_folding	True	能减少ONNX中的冗余计算节点
export_params	True	不导出权重的话，后面没法转TRT
keep_initializers_as_inputs	False	减少输入列表的冗余项

这里有个坑——opset_version不是越高越好。我踩过这个坑：用opset=15导出，TensorRT直接报不支持的算子。后来老老实实退回11，一切正常。工业场景求稳，别追新。

4.2 动态轴与静态轴

工业检测里，输入尺寸经常不固定。比如有的产线拍到的产品大小不一样，你不可能每次都resize到固定尺寸。这时候就需要动态轴。

动态轴说白了就是告诉ONNX：这个维度在推理时可以变化。最常见的动态轴是batch size和图像宽高。

dynamic_axes = {
    "input": {
        0: "batch_size",    # 动态batch
        2: "height",        # 动态高
        3: "width"          # 动态宽
    },
    "output": {
        0: "batch_size",
        2: "height",
        3: "width"
    }
}

注意：动态轴越多，TensorRT优化空间越小。我建议只对必要的维度设动态，其他保持静态。比如你的模型只接受640x640输入，那就只设batch为动态，宽高固定。

为什么会这样？因为TensorRT在做层融合时，需要知道每个张量的确切形状。动态轴意味着它要做更多运行时推断，优化效果自然打折扣。

我个人的经验是：能用静态就别用动态。如果产线图像尺寸变化不大，干脆统一resize到固定尺寸，性能能提升20%以上。

4.3 算子兼容性检查

这是最容易出问题的一环。PyTorch里很多操作在ONNX里没有对应算子，导出时会自动拆成多个基础算子。拆得好还行，拆不好就出幺蛾子。

常见的坑有哪些？

torch.where：ONNX里没有直接对应，会拆成Cast + NonZero + Gather，容易出精度问题
torch.topk：opset<11时支持不好，建议升到11以上
torch.nn.functional.grid_sample：工业检测里做仿射变换常用，ONNX支持但TensorRT支持有限
自定义算子：比如你写了个C++扩展，ONNX根本不认识

检查方法：导出后用onnxruntime跑一遍，对比PyTorch输出。如果误差超过1e-5，说明算子拆解有问题。

import onnxruntime as ort
import numpy as np

# 用ONNX Runtime验证
ort_session = ort.InferenceSession("model.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)

# 对比PyTorch输出
with torch.no_grad():
    torch_outputs = model(dummy_input)

# 计算误差
diff = np.abs(ort_outputs[0] - torch_outputs.numpy())
print(f"最大误差: {diff.max():.6f}")

我曾经遇到一个案例：模型里用了F.interpolate做上采样，默认的mode='nearest'在ONNX里没问题，但换成mode='bilinear'后，TensorRT推理结果和PyTorch差了0.5个像素。查了两天才发现是align_corners参数默认值不一致。PyTorch里默认False，但ONNX导出时某些版本会强制转成True。

4.4 导出流程与验证

我总结了一套标准流程，每次导出都按这个来，基本没出过问题。

模型准备：确保model.eval()，关闭dropout和batch norm的训练模式
dummy input构造：形状、数值范围要和真实数据一致
导出：按上面的参数设置，opset=11起步
ONNX检查：用onnx.checker.check_model验证结构完整性
精度验证：用onnxruntime跑推理，对比PyTorch输出
TensorRT初步转换：用trtexec转一次，看有没有报错

避坑指南：我曾经在导出时忘了设model.eval()，结果ONNX里包含了dropout层，推理结果每次都不一样。排查了整整一天才发现这个低级错误。

4.5 知识结构图

下面这张图概括了ONNX导出的核心流程和关键决策点，建议你保存下来，每次导出前对照检查。

4.6 常见问题与解决方案

问题现象	原因	解决方案
导出时报"Unsupported operator"	PyTorch算子没有ONNX映射	替换为ONNX支持的算子，或用torch.onnx.register_custom_op
ONNX Runtime推理结果全0	export_params=False	设置export_params=True
动态轴设置后TensorRT报错	某些层不支持动态形状	减少动态轴数量，或使用TRT的优化策略
精度误差超过1%	算子拆解导致数值差异	检查align_corners、eps等参数默认值