使用TensorRT进行模型benchmark的标准流程
在AI系统从实验室走向生产环境的过程中，一个常被忽视但至关重要的环节是：模型推理性能到底能不能扛住真实流量？
训练完成的模型精度再高，如果在线上服务中延迟飙升、吞吐不足，用户体验依然会崩塌。尤其是在视频分析、语音交互、自动驾驶这类对实时性要求极高的场景里，毫秒级的延迟差异可能直接决定产品成败。
NVIDIA TensorRT 正是在这个背景下成为工业部署的“性能加速器”。它不是训练框架，也不是新的神经网络结构，而是一套专为推理优化打造的工具链。它的核心使命很明确：把已经训练好的模型，在特定GPU上跑得更快、更省资源。
要真正发挥它的价值，不能靠“试试看”，而是需要一套标准、可复现的 benchmark 流程——这正是本文要讲清楚的事。
从ONNX到.engine：一次典型的优化之旅
假设你手头有一个PyTorch训练好的ResNet50模型，现在想看看它在T4服务器上的实际表现。第一步不是直接测速，而是准备好中间格式：
import torch import torchvision.models as models # 导出为ONNX model = models.resnet50(pretrained=True).eval() dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, "resnet50.onnx", input_names=["input"], output_names=["output"], opset_version=13, do_constant_folding=True, dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}} )这里有几个关键点容易踩坑：
-opset_version建议使用11及以上，太旧的版本可能不支持TensorRT的一些操作；
-dynamic_axes如果未来要支持变长batch或分辨率，必须提前声明；
- 输出的ONNX最好用onnx.checker.check_model()验证一下合法性。
一旦拿到ONNX文件，就进入了TensorRT的核心阶段——构建推理引擎。
import tensorrt as trt import pycuda.driver as cuda TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def build_engine(model_path, precision="fp16", workspace=1 << 30): builder = trt.Builder(TRT_LOGGER) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser = trt.OnnxParser(network, TRT_LOGGER) with open(model_path, 'rb') as f: if not parser.parse(f.read()): for i in range(parser.num_errors): print(parser.get_error(i)) raise RuntimeError("ONNX parsing failed") config = builder.create_builder_config() config.max_workspace_size = workspace # 单位字节 if precision == "fp16" and builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) elif precision == "int8": config.set_flag(trt.BuilderFlag.INT8) # 这里需要实现校准器（后文详述） return builder.build_serialized_network(network, config)这段代码看似简单，实则藏着不少门道。比如max_workspace_size设置得太小，可能导致某些层无法启用最优算法；设置得太大又浪费显存。经验上，1GB（1<<30）对于大多数CV模型够用，但像BERT-Large这样的大模型可能需要4~8GB。
最终生成的是一个.engine文件，它是完全序列化的、平台相关的二进制体，包含了所有优化后的计算图和内核选择。这意味着你可以在没有原始框架依赖的环境中加载它，真正做到“一次编译，随处运行”（当然前提是同架构GPU）。
benchmark不只是跑个平均延迟
很多人理解的benchmark就是“喂数据、计时、取平均”，但这远远不够。真正的性能评估应该回答以下几个问题：
- 在不同batch size下，QPS如何变化？
- GPU利用率是否饱和？有没有空转？
- 显存占用会不会成为瓶颈？
- FP16真的没掉点吗？INT8还能接受吗？
为此，我们需要一个完整的推理执行循环：
import numpy as np class EngineRunner: def __init__(self, engine_data): self.runtime = trt.Runtime(TRT_LOGGER) self.engine = self.runtime.deserialize_cuda_engine(engine_data) self.context = self.engine.create_execution_context() # 分配host/device缓冲区 self.inputs = [] self.outputs = [] self.bindings = [] for i in range(self.engine.num_bindings): binding = self.engine[i] size = trt.volume(self.engine.get_binding_shape(binding)) * self.engine.max_batch_size dtype = trt.nptype(self.engine.get_binding_dtype(binding)) host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(device_mem.nbytes) self.bindings.append(int(device_mem)) if self.engine.binding_is_input(binding): self.inputs.append({'host': host_mem, 'device': device_mem}) else: self.outputs.append({'host': host_mem, 'device': device_mem}) def infer(self, input_data, stream): # H2D np.copyto(self.inputs[0]['host'], input_data.ravel()) cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], stream) # 执行 self.context.execute_async_v2(bindings=self.bindings, stream_handle=stream.handle) # D2H cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], stream) stream.synchronize() return self.outputs[0]['host'].reshape(self.engine.get_binding_shape(1))注意这里用了execute_async_v2和 CUDA Stream 实现异步传输与计算重叠，这才是贴近真实服务场景的做法。同步调用execute()虽然简单，但会人为拉高延迟。
正式测试前别忘了预热（warm-up）：
# Warm-up stream = cuda.Stream() dummy_input = np.random.rand(1, 3, 224, 224).astype(np.float32) for _ in range(10): runner.infer(dummy_input, stream)GPU有很多缓存机制（L2 cache、TLB等），冷启动时性能往往偏低。预热能让硬件进入稳定状态，使后续测量更准确。
性能指标怎么才算“好”？
我们通常关注三个核心指标：
| 指标 | 定义 | 测量方式 |
|---|---|---|
| 延迟（Latency） | 单次请求从输入到输出的时间 | 取多次运行的均值/99分位数 |
| 吞吐量（QPS） | 每秒能处理的请求数 | 总请求数 / 总耗时 |
| GPU利用率 | SM活跃时间占比 | nvidia-smi dmon或 Nsight Systems |
举个例子，在T4上运行ResNet50：
| 精度模式 | Batch=1延迟 | Batch=32 QPS | 相比原生PyTorch提升 |
|---|---|---|---|
| FP32 | ~12ms | ~2800 | - |
| FP16 | ~6.5ms | ~5000 | 吞吐+78% |
| INT8 | ~3.8ms | ~8600 | 吞吐+207% |
可以看到，仅通过FP16就能实现接近翻倍的性能收益，而INT8更是将吞吐推到了极致。不过代价也很明显：INT8需要额外的校准步骤，且并非所有模型都能承受这种量化带来的精度损失。
说到INT8，很多人失败的原因在于随便选几条数据做校准。正确的做法是使用具有代表性的数据集（例如ImageNet validation set的子集），并采用合适的校准策略：
- Entropy Calibration：基于KL散度最小化，适合大多数图像分类任务；
- MinMax Calibration：取激活值的最大最小范围，简单但易受异常值影响；
- Percentile Calibration：忽略极端百分位的值，鲁棒性更强。
class Int8Calibrator(trt.IInt8EntropyCalibrator2): def __init__(self, data_loader, cache_file): super().__init__() self.data_loader = data_loader self.dummy_input = cuda.mem_alloc(data_loader.batch_size * 3 * 224 * 224 * 4) self.cache_file = cache_file self.current_batch_idx = 0 def get_batch(self, names): if self.current_batch_idx >= len(self.data_loader): return None batch = next(iter(self.data_loader)) np.copyto(self.host_mem, batch.numpy().ravel()) cuda.memcpy_htod(self.dummy_input, self.host_mem) self.current_batch_idx += 1 return [int(self.dummy_input)] def read_calibration_cache(self): pass def write_calibration_cache(self, cache): with open(self.cache_file, 'wb') as f: f.write(cache)校准样本数量建议控制在100~500之间。太少会导致统计不准，太多则增加构建时间且边际效益递减。
实战中的常见陷阱与应对
1. “我的ONNX模型解析失败！”
最常见的原因是操作符不支持或动态shape未正确配置。解决方案：
- 查阅TensorRT官方支持的操作列表；
- 尝试用polygraphy工具定位具体哪一层出错；
- 对于自定义算子，考虑用Plugin机制扩展。
2. “为什么batch增大后QPS反而下降？”
这通常是显存带宽成了瓶颈，或者GPU SM利用率未饱和。可通过以下方式排查：
- 使用Nsight Systems分析kernel执行间隔是否有空洞；
- 检查是否启用了层融合（查看profile日志）；
- 尝试调整builder_config.set_tactic_sources()强制启用更多优化策略。
3. “INT8之后top-1精度掉了5个点怎么办？”
不要轻易放弃INT8，先检查：
- 是否使用了正确的校准数据分布；
- 是否在敏感层（如最后的全连接层）禁用了量化；
- 是否可以结合混合精度（部分层保持FP16）。
有时候，微调校准参数比重新训练模型成本更低。
如何让benchmark变成自动化流水线？
在团队协作中，手动跑脚本显然不可持续。理想的做法是将整个流程接入CI/CD：
# .github/workflows/benchmark.yml name: Model Benchmark on: [push] jobs: trt-benchmark: runs-on: ubuntu-latest container: nvcr.io/nvidia/tensorrt:23.09-py3 steps: - name: Checkout uses: actions/checkout@v3 - name: Export ONNX run: python export_onnx.py - name: Build & Test Engine run: | python build_engine.py --precision fp16 python benchmark.py --engine resnet50_fp16.engine --batch-sizes 1,8,16,32 - name: Upload Results run: curl -X POST $PERF_DASHBOARD_ENDPOINT --data @results.json每次模型更新，自动触发benchmark，对比历史数据判断是否存在性能回归。这种闭环机制能极大提升迭代效率。
写在最后：benchmark的本质是建立信任
使用TensorRT进行benchmark，表面上是在测速度，实质上是在回答一个问题：这个模型上线后，我能睡安稳觉吗？
它迫使我们跳出“训练即终点”的思维定式，转而思考部署侧的真实约束——延迟、吞吐、资源占用、稳定性。而TensorRT提供的，正是一套系统化的方法论，让我们能把模糊的“感觉还行”转化为清晰的数字证据。
更重要的是，这套流程本身是可以沉淀的。当你为第一个模型走通了从ONNX导出到INT8校准的全链路，后续项目就能快速复制。这种工程能力的积累，才是AI落地中最宝贵的资产。
所以，下次拿到一个新模型时，不妨先别急着吹嘘它的精度有多高，而是问一句：它在T4上跑多少毫秒？
