3. ONNX模型导出与检查

ONNX（Open Neural Network Exchange）是TensorRT部署链路中最常见的中间表示。TensorRT本身并不直接读取PyTorch的.pt或TensorFlow的SavedModel，而是通过ONNX Parser把计算图翻译成TensorRT的INetwork定义。因此，ONNX文件的质量几乎决定了后续INT8校准、层融合和动态Shape支持的成败。本节围绕“导出—检查—简化—验证”四个动作展开，重点讲清楚每一步实际发生了什么、参数如何影响结果、以及如何判断导出的模型是否可用。

一、导出前的准备：把模型切到推理态

无论PyTorch还是TensorFlow，导出前必须让模型进入推理模式。PyTorch需要调用model.eval()，否则Dropout和BatchNorm会保留训练行为，导出的ONNX会带有随机的Dropout节点或错误的BN统计量。同时建议用torch.no_grad()包裹导出过程，避免导出时构建反向图。TensorFlow则使用tf.function配合input_signature固定输入签名，再通过tf2onnx转换。

PyTorch导出ONNX的核心接口是torch.onnx.export，关键参数包括：opset_version、input_names、output_names、dynamic_axes、do_constant_folding。下面是一个典型的导出示例，模型是一个简单的卷积分类网络，输入为[1,3,224,224]：

import torch
import torch.onnx

model = MyNet().eval()
dummy = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy,
    "mynet.onnx",
    opset_version=13,
    input_names=["input"],
    output_names=["logits"],
    dynamic_axes={
        "input":  {0: "batch", 2: "height", 3: "width"},
        "logits": {0: "batch"}
    },
    do_constant_folding=True
)

二、opset选择：不是越高越好

opset_version决定了ONNX算子集的版本。版本越高，支持的算子越丰富，但TensorRT的ONNX Parser对高版本opset的支持存在滞后。以TensorRT 8.x为例，官方推荐使用opset 11～13；opset 17及以上虽然能被解析，但部分新算子（如LayerNormalization、GroupNormalization的某些变体）可能回退到插件或直接报错。选择原则是：优先使用模型导出时框架默认支持的、且TensorRT明确支持的opset。如果模型包含Transformer结构，opset 13通常能覆盖MultiHeadAttention的分解形式；如果包含GridSample等算子，需要确认TensorRT版本是否支持。

三、动态轴设置：让batch和空间维度可变

动态轴是TensorRT部署中非常关键的一环。如果导出时把batch固定为1，后续想支持batch=8就必须重新导出。动态轴通过dynamic_axes参数声明，格式为{张量名: {维度索引: 维度名称}}。维度名称只是标记，真正起作用的是“该维度被声明为动态”。

需要注意三点：第一，动态轴必须从输入和输出两端同时声明，否则ONNX检查会报维度不匹配；第二，并非所有维度都适合动态，例如分类模型的类别数维度通常固定；第三，动态Shape会显著增加TensorRT的优化难度，因为Builder需要为每个动态维度生成多个优化配置（Optimization Profile）。在TensorRT中，动态维度需要指定min/opt/max三组值，例如batch维度设为[1,8,16]，表示最小1、最优8、最大16。

四、模型简化与onnxruntime验证

导出的ONNX往往包含冗余节点，例如恒等变换、多余的Transpose、可折叠的常量运算。这些节点会干扰TensorRT的层融合，甚至导致某些子图无法被识别。常用工具是onnx-simplifier，它通过常量折叠、算子消除和形状推断来精简图结构：

python -m onnxsim mynet.onnx mynet_sim.onnx

简化后必须用onnxruntime做数值验证，确认简化没有改变计算结果。验证方法是：用同一组输入分别跑原始ONNX和简化ONNX，比较输出差异。如果最大绝对误差在1e-5以内，可以认为简化是安全的。onnxruntime还能帮助发现算子不支持、输入输出名称错误、动态轴声明不一致等问题。

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("mynet_sim.onnx")
x = np.random.randn(1, 3, 224, 224).astype(np.float32)
out = sess.run(["logits"], {"input": x})
print(out[0].shape)   # 期望 (1, 1000)

如果onnxruntime报错“Invalid graph: Dynamic shape mismatch”，通常说明dynamic_axes声明与算子实际推导的形状冲突，需要回到导出步骤检查。如果报错“Unsupported operator”，则说明opset或算子版本超出onnxruntime支持范围，需要调整opset或替换算子实现。

五、常见问题与判断依据

导出过程中最常遇到三类问题。第一类是形状不匹配：onnxruntime加载时报“Got invalid dimensions for input”，通常是因为dynamic_axes中声明的维度名称与输入张量实际维度不一致，或者输出端漏声明。第二类是算子不支持：TensorRT Parser报“Unsupported ONNX op: XXX”，需要查TensorRT版本对应的ONNX算子支持列表，必要时用自定义插件替代。第三类是精度异常：简化后输出误差超过阈值，说明onnxsim的常量折叠改变了数值行为，此时应关闭简化或改用--skip-fuse-bn等保守选项。

一个实用的检查顺序是：先用onnxruntime确认模型能跑通且输出形状正确；再用onnx.checker.check_model做结构校验；最后用polygraphy或TensorRT的trtexec做解析测试。如果trtexec --onnx=mynet_sim.onnx --verbose能成功构建引擎，说明ONNX已经具备进入TensorRT部署链路的基本条件。

总结来说，ONNX导出不是“一键完成”的动作，而是一个需要反复检查的工程环节。opset决定算子可用性，dynamic_axes决定部署灵活性，onnxsim决定图结构是否干净，onnxruntime决定数值是否正确。四者缺一不可，任何一步跳过都可能在TensorRT构建阶段暴露问题，而那时排查成本会显著上升。