TensorRT性能优化：如何提升Python推理脚本的执行效率

问题背景

在使用TensorRT进行模型推理时，开发者经常会遇到一个常见问题：使用trtexec工具转换和运行模型时获得的推理速度，往往比使用自定义Python脚本实现的推理速度快很多。本文将以一个YOLO模型为例，分析这种性能差异的原因，并提供优化建议。

性能差异分析

在案例中，开发者发现：

• 使用trtexec工具时，推理时间在3-5毫秒范围内
• 使用自定义Python脚本时，推理时间增加到10-12毫秒

这种性能差距主要来自以下几个方面：

1. CUDA初始化开销

Python脚本首次运行时，CUDA环境需要初始化，这会带来额外的开销。专业的性能测试应该包含"预热"阶段，即在正式测量前先运行几次推理，使CUDA环境达到稳定状态。

2. 时间测量方式不准确

原脚本测量的是端到端(E2E)时间，包含了：

• 主机到设备(H2D)数据传输时间
• GPU计算时间
• 设备到主机(D2H)数据传输时间
• Python环境开销

而trtexec工具会分别测量这些时间组件，提供更精确的性能分析。

3. 执行上下文创建

每次推理都创建新的执行上下文会增加额外开销。最佳实践是在预热阶段创建并复用执行上下文。

优化建议

1. 添加CUDA预热

在正式测量前，先运行几次推理任务：

# 预热阶段
for _ in range(10):  # 预热10次
    _, _ = Inference(engine, warmup_image_path)

# 正式测量
execution_times = []
for image_path in inf_images:
    output, execution_time = Inference(engine, image_path)
    execution_times.append(execution_time)

2. 精确时间测量

区分不同阶段的时间测量：

start_enqueue = time.time()
cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
enqueue_time = time.time() - start_enqueue

start_compute = time.time()
context.execute_v2(bindings)
compute_time = time.time() - start_compute

start_d2h = time.time()
cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
stream.synchronize()
d2h_time = time.time() - start_d2h

total_time = enqueue_time + compute_time + d2h_time

3. 执行上下文复用

避免在每次推理时创建新的执行上下文：

# 在PrepareEngine中创建并保存上下文
context = engine.create_execution_context()

# 在Inference函数中直接使用已创建的上下文
def Inference(context, image_path):
    # ...其他代码不变...
    context.execute_v2(bindings)
    # ...其他代码不变...

4. 批处理优化

如果可能，使用更大的批处理大小来提高GPU利用率：

# 修改PrepareEngine中的批处理大小
batch = 4  # 根据GPU内存调整

性能指标解读

理解TensorRT的性能指标对于优化至关重要：

1. 总主机墙钟时间：从第一个查询入队到最后一个查询完成的全部时间
2. GPU计算时间：GPU执行内核的实际计算时间
3. 总GPU计算时间：所有查询GPU计算时间的总和
4. 吞吐量：查询数量除以总主机墙钟时间
5. 入队时间：主机将查询入队的延迟
6. H2D延迟：输入张量从主机到设备的数据传输时间
7. D2H延迟：输出张量从设备到主机的数据传输时间
8. 延迟：H2D延迟、GPU计算时间和D2H延迟的总和

结论

通过上述优化措施，Python推理脚本的性能可以接近trtexec工具的水平。关键在于：

• 正确的预热过程
• 精确的时间测量和分析
• 合理的资源复用
• 适当的批处理设置

理解TensorRT的工作原理和性能特征，能够帮助开发者编写出更高效的推理代码，充分发挥硬件加速的潜力。