4、ONNX模型调试与验证：ONNX Runtime推理、onnxruntime-gpu与onnxruntime-silicon、ONNX模型正确性验证

模型导出成 ONNX 之后，别急着往 TensorRT 里塞。我见过太多人，导出来就直接转，结果跑出来的结果全是错的，还以为是 TensorRT 的锅。

其实，ONNX 这一步是个很好的「中间检查点」。你得先确认 ONNX 本身是对的，再去谈优化。说白了，这一步就是给模型做个「体检」。

4.1 为什么要在 ONNX 阶段做验证？

你想想看，从 PyTorch 到 TensorRT，中间隔了好几层转换。如果 ONNX 这一步就有问题，后面排查起来会非常痛苦。

我个人习惯是，导出 ONNX 后，先用 ONNX Runtime 跑一遍推理，跟原始框架（比如 PyTorch）的输出做对比。误差在 1e-5 以内，我才敢往下走。

4.2 ONNX Runtime 推理：CPU 下的快速验证

ONNX Runtime（简称 ORT）是微软开源的推理引擎。它最大的好处是：不需要 GPU，装个 pip 包就能跑。非常适合做快速验证。

4.2.1 安装与基础用法

pip install onnxruntime

装完之后，写个简单的推理脚本：

import onnxruntime as ort
import numpy as np

# 创建推理会话
session = ort.InferenceSession("model.onnx")

# 获取输入输出信息
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 构造随机输入（注意 shape 和 dtype 要匹配）
dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 执行推理
outputs = session.run([output_name], {input_name: dummy_input})

print(outputs[0].shape)

嗯，这里要注意：session.run 的第一个参数是输出名称列表，传 None 会返回所有输出。但我建议显式指定，避免不必要的计算。

4.2.2 数值比对脚本

我一般会写一个比对函数，专门用来验证 ONNX 和原始框架的输出：

def compare_outputs(torch_output, onnx_output, rtol=1e-5, atol=1e-5):
    """
    比对 PyTorch 和 ONNX Runtime 的输出
    """
    torch_np = torch_output.detach().cpu().numpy()
    onnx_np = onnx_output

    max_diff = np.max(np.abs(torch_np - onnx_np))
    mean_diff = np.mean(np.abs(torch_np - onnx_np))

    print(f"最大误差: {max_diff:.2e}")
    print(f"平均误差: {mean_diff:.2e}")

    if np.allclose(torch_np, onnx_np, rtol=rtol, atol=atol):
        print("✅ 验证通过！")
    else:
        print("❌ 验证失败！")

    return max_diff, mean_diff

4.3 onnxruntime-gpu 与 onnxruntime-silicon

CPU 验证通过后，你可能想看看 ONNX 在 GPU 上的表现。这时候就需要 onnxruntime-gpu 了。

4.3.1 版本选择

安装 GPU 版本时，版本号一定要跟你的 CUDA 版本匹配。我曾经因为装了不匹配的版本，折腾了一下午才发现是版本问题。

# CUDA 11.x 用户
pip install onnxruntime-gpu

# CUDA 12.x 用户
pip install onnxruntime-gpu==1.17.0  # 具体版本请查官方文档

4.3.2 GPU 推理代码

跟 CPU 版本几乎一样，只是创建 session 时指定 provider：

import onnxruntime as ort

# 创建 GPU 会话
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
session = ort.InferenceSession("model.onnx", providers=providers)

# 后面的推理代码跟 CPU 版本完全一致
# ...

4.4 ONNX 模型正确性验证

验证不只是跑一次推理那么简单。我总结了一套「三步验证法」：

1. 结构验证： 用 onnx.checker 检查模型结构是否合法。
2. 数值验证： 用多组输入比对输出，覆盖不同数据分布。
3. 动态形状验证： 如果模型支持动态 batch，要测试不同 batch size 下的正确性。

4.4.1 结构验证

import onnx

# 加载模型
model = onnx.load("model.onnx")

# 检查模型结构
onnx.checker.check_model(model)

print("✅ 模型结构检查通过")

这一步能发现很多低级错误，比如某个节点的输入输出名称不匹配、shape 信息缺失等。

4.4.2 多组输入验证

只用一组随机输入验证是不够的。我建议至少准备 3-5 组不同的输入：

• 全零输入： 检查模型在极端情况下的行为。
• 全一输入： 同上。
• 随机输入： 模拟正常数据分布。
• 真实数据： 从训练集里拿几张图，用预处理后的 tensor 做输入。

test_inputs = [
    np.zeros((1, 3, 224, 224), dtype=np.float32),
    np.ones((1, 3, 224, 224), dtype=np.float32),
    np.random.randn(1, 3, 224, 224).astype(np.float32),
    # 真实预处理后的数据
]

for i, test_input in enumerate(test_inputs):
    onnx_output = session.run([output_name], {input_name: test_input})
    # 跟原始框架输出做比对
    # ...

4.4.3 动态形状验证

如果你的 ONNX 模型支持动态 batch，一定要测试不同 batch size：

batch_sizes = [1, 2, 4, 8]

for bs in batch_sizes:
    dummy_input = np.random.randn(bs, 3, 224, 224).astype(np.float32)
    outputs = session.run([output_name], {input_name: dummy_input})
    assert outputs[0].shape[0] == bs, f"Batch size {bs} 验证失败"
    print(f"Batch size {bs}: ✅")

4.5 知识结构图：ONNX 验证流程

下面这张图总结了我平时做 ONNX 验证的完整流程：

4.6 常见问题与排查思路

验证过程中，你可能会遇到各种问题。我整理了几个高频场景：

好了，ONNX 验证这一步走完，你的模型就已经通过了「体检」。接下来，就可以放心地把它交给 TensorRT 去做优化了。