第四章：基准模型构建
说实话，很多做TensorRT部署的朋友，一上来就急着调精度、搞优化。结果折腾半天，发现模型压根就没跑对。嗯，我当年也犯过这个错。
基准模型构建，说白了就是给你的模型画一条「及格线」。你得先确认原始模型是对的，后面做的一切优化才有意义。这一章，我们就来聊聊怎么搭好这条线。
4.1 从PyTorch/TensorFlow导出标准ONNX
我个人习惯，不管用什么框架训练，最终都会转成ONNX。为什么？因为ONNX是中间格式，TensorRT吃得最香。
先看PyTorch怎么导出。代码其实很简单：
import torch
import torch.onnx
# 假设你有一个训练好的模型
model = torch.load('my_model.pth')
model.eval()
# 构造一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)
# 导出 ONNX
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。太高了，有些老算子不支持；太低了，新特性用不了。
TensorFlow那边稍微麻烦点。我记得有一次，客户给的是TF1.x的模型，折腾了我一整天：
import tensorflow as tf
import tf2onnx
# TF2 导出
model = tf.saved_model.load('saved_model_dir')
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = model.signatures['serving_default']
# 转 ONNX
onnx_model, _ = tf2onnx.convert.from_function(
output_path,
input_signature=spec,
opset=13,
output_path="model.onnx"
)
TF1.x的话，你得用 freeze_graph 先把图冻结，再转。嗯，这里就不展开了，遇到再说。
4.2 ONNX算子兼容性检查
模型导出来了，你以为就完事了？别急。ONNX模型里可能藏着一些「定时炸弹」——不兼容的算子。
我常用的检查工具是 onnxruntime 自带的检查器，还有 onnx 库里的 checker：
import onnx
# 加载模型
model = onnx.load('model.onnx')
# 检查模型结构
onnx.checker.check_model(model)
print("模型结构检查通过")
# 查看所有算子
ops = set()
for node in model.graph.node:
ops.add(node.op_type)
print("模型中使用的算子:", ops)
如果检查报错，别慌。常见的问题就几种：
| 问题类型 | 表现 | 解决办法 |
|---|---|---|
| 算子版本不匹配 | checker 报版本错误 | 升级 opset 或降级算子 |
| 自定义算子 | 模型里有非标准算子 | 用 onnx.register_custom_op 注册 |
| 动态形状问题 | 某些算子不支持动态轴 | 固定形状或用 Reshape 替代 |
4.3 使用ONNX Runtime验证正确性
模型结构没问题了，接下来得验证精度。说白了，就是拿ONNX跑一遍，跟原始框架的结果对比。
我一般这么干：
import onnxruntime as ort
import numpy as np
# 创建推理会话
session = ort.InferenceSession('model.onnx')
# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 运行推理
outputs = session.run(
['output'],
{'input': input_data}
)
print("ONNX Runtime 输出形状:", outputs[0].shape)
print("ONNX Runtime 输出前5个值:", outputs[0][0][:5])
然后跟PyTorch的结果对比：
# PyTorch 推理
with torch.no_grad():
torch_output = model(torch.from_numpy(input_data))
# 计算误差
diff = np.abs(
outputs[0] - torch_output.numpy()
)
max_diff = np.max(diff)
mean_diff = np.mean(diff)
print(f"最大误差: {max_diff:.6f}")
print(f"平均误差: {mean_diff:.6f}")
# 判断是否通过
if max_diff < 1e-5:
print("✅ 精度验证通过")
else:
print("❌ 精度验证失败，需要排查")
你可能会问，为什么要用 ONNX Runtime 验证？直接用 TensorRT 不行吗？
嗯，这里有个关键点。ONNX Runtime 是 ONNX 的官方推理引擎，它跑出来的结果可以看作是「标准答案」。如果 ONNX Runtime 的结果跟原始框架不一致，那说明导出过程就有问题。这时候你调 TensorRT 的精度，根本调不明白。
我曾经帮一个客户排查问题，他死活说 TensorRT 精度不对。我让他先跑 ONNX Runtime，结果发现 ONNX Runtime 跟 PyTorch 的误差已经到 0.1 了。最后查出来，是导出时有个参数没设对。所以，基准验证这一步，千万别跳。
4.4 常见问题与避坑指南
最后，分享几个我实际项目中遇到的坑：
- 动态轴设置不当：导出时 dynamic_axes 设得太随意，导致某些算子不支持。我建议只对 batch 维度设动态，其他维度尽量固定。
- 数据类型不一致：PyTorch 默认 float32，但有些模型用了 float16 或 double。导出前统一转成 float32。
- 控制流问题：模型里有 if-else 或循环，导出后可能变成一堆奇怪的算子。尽量用 torch.where 或 masked 操作替代。
下一章，我们会聊到 ONNX 模型的可视化与调试。到时候，我会教你用 Netron 看模型结构，用 onnx_graphsurgeon 做精细调整。嗯，那些都是实战中特别有用的技能。