4. ONNX模型获取与检查

在TensorRT的模型转换链路中，ONNX（Open Neural Network Exchange）扮演着“通用中间语言”的角色。无论是PyTorch、TensorFlow还是其他主流框架训练的模型，最终都需要统一转换为ONNX格式，才能被TensorRT的解析器（onnx-tensorrt或trtexec）高效消费。本章聚焦于ONNX模型获取与检查的完整闭环：从源框架导出、到可视化结构审查、再到运行时数值验证。这一环节的质量直接决定了后续TensorRT引擎构建的成败——如果ONNX本身存在算子不兼容、动态维度定义错误或精度损失，那么无论TensorRT优化多出色，最终推理结果都可能是错误的。

4.1 核心原理：为什么需要ONNX作为中间表示

ONNX本质上是一种基于计算图的协议格式，它将神经网络描述为一系列算子（Operator）节点和张量（Tensor）边的有向无环图（DAG）。与PyTorch的即时执行（Eager Execution）或TensorFlow的图执行模式不同，ONNX提供了一种框架无关的静态图表示。这意味着模型的结构、权重、数据流和维度信息都被显式地序列化到文件中，不依赖任何特定框架的运行时环境。

从工程角度看，ONNX的价值体现在三个层面：第一，解耦训练与部署——训练阶段可以使用PyTorch的灵活性，部署阶段则可以切换到TensorRT的高性能推理引擎，两者通过ONNX无缝衔接；第二，算子标准化——ONNX定义了约150个标准算子（如Conv、MatMul、Gemm、Reshape等），这些算子在各个推理后端都有对应的实现，避免了直接对接框架私有算子的复杂性；第三，优化前置——在ONNX阶段就可以进行图优化（如常量折叠、算子融合），减少后续TensorRT解析的负担。

然而，ONNX并非万能。它存在一个关键约束：算子集版本（Opset Version）。每个ONNX算子都有其引入和演进的版本，导出时指定的opset_version决定了可用的算子集合。如果opset版本过低，某些新算子无法表达；如果过高，TensorRT解析器可能尚未支持。因此，选择合适的opset版本是导出阶段的首要决策。

4.2 从PyTorch导出ONNX：操作流程与关键参数

PyTorch提供了torch.onnx.export接口，其核心机制是通过追踪（Tracing）或脚本化（Scripting）方式记录模型的前向计算过程。对于大多数CNN模型，追踪方式足够；对于包含数据依赖控制流的模型（如循环、条件分支），则需要使用Scripting或混合模式。

以下是一个标准的导出流程示例：

import torch
import torchvision.models as models

# 1. 加载模型并设置为评估模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造示例输入（必须与模型输入维度匹配）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 定义动态轴（可选但推荐）
dynamic_axes = {
    "input": {0: "batch_size"},   # 第0维动态
    "output": {0: "batch_size"}
}

# 4. 执行导出
torch.onnx.export(
    model,                          # 模型
    dummy_input,                    # 示例输入
    "resnet18.onnx",                # 输出文件路径
    export_params=True,             # 导出权重
    opset_version=17,               # 算子集版本
    do_constant_folding=True,       # 常量折叠优化
    input_names=["input"],          # 输入张量名
    output_names=["output"],        # 输出张量名
    dynamic_axes=dynamic_axes       # 动态维度定义
)
print("ONNX导出完成")

导出过程中有几个容易踩坑的细节：第一，BatchNorm层必须固定——在导出前调用model.eval()，否则BatchNorm会使用批统计量而非全局统计量，导致导出的权重错误；第二，示例输入的数据类型——必须与训练时一致，通常为float32；第三，动态维度的粒度——可以只动态batch维度，也可以动态宽高（如检测模型），但动态维度越多，TensorRT优化空间越小。

4.3 从TensorFlow导出ONNX：使用tf2onnx工具链

TensorFlow模型导出ONNX通常借助tf2onnx库。与PyTorch不同，TensorFlow有多个模型保存格式（SavedModel、Keras H5、Checkpoint），tf2onnx主要支持从SavedModel或Keras模型直接转换。

# 方式一：命令行转换（推荐）
python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 17 \
    --inputs input_tensor:0 \
    --outputs output_tensor:0

# 方式二：Python API转换
import tf2onnx
import tensorflow as tf

model = tf.keras.models.load_model("my_model.h5")
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = "./model.onnx"
model.output_names = ["output"]

onnx_model = tf2onnx.convert.from_keras(model, input_signature=spec, opset=17)
with open(output_path, "wb") as f:
    f.write(onnx_model.SerializeToString())

TensorFlow转换的常见问题包括：输入占位符命名——必须明确指定输入张量的名称，否则tf2onnx会生成随机名称；控制流算子——TensorFlow的tf.while_loop或tf.cond会转换为ONNX的Loop或If算子，这些在TensorRT中支持有限；NHWC与NCHW布局——TensorFlow默认使用NHWC（通道在后），而TensorRT和PyTorch通常使用NCHW，需要在转换时通过--inputs-as-nchw参数进行布局转换。

4.4 使用Netron进行可视化检查：结构审查方法论

Netron是微软开源的可视化工具，支持打开ONNX、TensorFlow、PyTorch等多种格式。它不仅是“看图工具”，更是结构审查的利器。打开导出的ONNX文件后，应按照以下顺序进行检查：

1. 输入输出节点检查：点击输入节点，确认张量名称、数据类型（应为float32）、形状（动态维度显示为?或batch_size）。输出节点同样需要确认。
2. 算子序列审查：沿着主路径检查算子类型是否符合预期。例如，ResNet应包含Conv、BatchNormalization、Relu、MaxPool、GlobalAveragePool、Gemm等。如果出现不期望的算子（如自定义Op），需要标记风险。
3. 权重张量验证：点击Conv节点，查看权重形状。例如，Conv2d的权重应为[out_channels, in_channels, kH, kW]，如果形状异常，说明导出时维度顺序出错。
4. 动态维度标记：在Netron中，动态维度会显示为?或自定义符号。如果模型本应支持动态batch，但显示为固定值，说明dynamic_axes未生效。
5. 子图结构：对于复杂模型（如检测模型），检查是否有子图（Subgraph）结构，这些通常对应TensorFlow的cond/while或PyTorch的torch.where，需要额外关注TensorRT兼容性。

4.5 onnxruntime验证：数值一致性与精度评估

可视化只能确认结构正确，无法验证数值正确性。因此，必须使用onnxruntime（简称ORT）进行运行时验证。验证的核心目标是：对于相同的输入，ONNX模型的输出与源框架（PyTorch/TensorFlow）的输出在数值上高度一致。

import onnxruntime as ort
import numpy as np
import torch

# 1. 创建推理会话
sess = ort.InferenceSession("resnet18.onnx", providers=["CPUExecutionProvider"])

# 2. 准备输入数据（与导出时的示例输入同分布）
np_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 3. 运行ONNX推理
ort_output = sess.run(None, {"input": np_input})[0]

# 4. 运行PyTorch推理（作为基准）
with torch.no_grad():
    torch_output = model(torch.from_numpy(np_input)).numpy()

# 5. 计算误差指标
abs_diff = np.abs(ort_output - torch_output)
cos_sim = np.dot(ort_output.flatten(), torch_output.flatten()) / (
    np.linalg.norm(ort_output) * np.linalg.norm(torch_output)
)
print(f"最大误差: {abs_diff.max():.6f}")
print(f"平均误差: {abs_diff.mean():.6f}")
print(f"余弦相似度: {cos_sim:.8f}")

# 6. 断言验证
assert abs_diff.max() < 1e-4, "误差过大，检查导出配置"
assert cos_sim > 0.9999, "余弦相似度不足"

验证过程中需要关注的指标包括：最大误差（应小于1e-4）、平均误差（应小于1e-5）、余弦相似度（应大于0.9999）。如果误差超出范围，可能的原因包括：BatchNorm统计量未固定、权重精度意外降低（如某些层被转换为float16）、算子实现差异（如Resize的插值算法不同）。

4.6 案例拆解：YOLOv5s模型导出与验证全流程

以YOLOv5s检测模型为例，展示完整的ONNX获取与检查闭环。该模型包含C3模块、SPPF、Detect头等复杂结构，是验证导出流程的理想案例。

背景与约束：项目需要将YOLOv5s部署到TensorRT，输入尺寸为640×640，batch size在1~8之间动态变化。初始现象是直接使用官方导出脚本生成的ONNX在TensorRT中构建失败，报错“Unsupported Operator: GridSample”。

分析判断过程：通过Netron打开ONNX文件，发现Detect头中使用了GridSample算子（用于坐标网格采样），该算子在TensorRT 8.x中不支持。进一步检查发现，YOLOv5的导出脚本默认启用了--simplify选项，但简化过程并未消除GridSample。

实际操作步骤：

1. 修改导出脚本，添加--dynamic参数启用动态batch，并指定--opset 17。
2. 导出后，使用Netron检查Detect头，确认GridSample算子存在。
3. 使用onnxruntime验证数值一致性，发现输出与PyTorch原始模型完全一致（误差<1e-6）。
4. 针对TensorRT兼容性，采用算子替换策略：将GridSample替换为等价的Resize + Gather组合，或使用--trt-extra-plugins加载自定义插件。
5. 最终选择使用TensorRT官方提供的EfficientNMS插件替换Detect头，彻底移除GridSample。

结果对比：替换后ONNX模型在TensorRT中构建成功，推理速度从原始PyTorch的12ms提升至2.1ms（FP16），mAP@0.5仅下降0.3%。

风险与代价：算子替换需要重新验证精度，且EfficientNMS的置信度阈值和NMS参数需要与原始实现对齐，否则可能影响检测效果。

4.7 常见误区与总结复盘

总结而言，ONNX获取与检查是TensorRT部署链路中的“质量门禁”。一个经过严格验证的ONNX模型，能够显著降低后续TensorRT构建和调优的复杂度。建议在项目初期就建立标准化的导出与验证脚本，将检查流程自动化，确保每次模型更新都能快速通过验证。记住：在ONNX阶段多花10分钟检查，可能节省TensorRT阶段数小时的排障时间。