TensorRT 8.6.1在NVIDIA GeForce RTX 3060上构建引擎时的内存泄漏问题分析

问题背景

在使用TensorRT 8.6.1版本将ONNX模型转换为TensorRT引擎时，开发者遇到了一个严重的内存泄漏问题。当在NVIDIA GeForce RTX 3060显卡上运行tensorrt.Builder.build_serialized_network方法时，系统内存会逐渐被耗尽，最终导致进程崩溃。这个问题在64GB内存的机器上需要数小时才会显现。

模型结构分析

出现问题的模型是一个简单的遍历性计算模型，主要包含以下关键操作：

1. 使用torch.where进行条件选择
2. 通过torch.minimum进行最小值计算
3. 包含两个嵌套循环结构，循环次数为MAX_TRAV*2+1
4. 使用torch.cat进行张量拼接

模型的核心算法是通过迭代计算来传播障碍物信息，最终输出每个位置到最近障碍物的距离。这种算法在机器人路径规划和计算机视觉中较为常见。

问题复现环境

• TensorRT版本：8.6.1
• GPU型号：NVIDIA GeForce RTX 3060
• 驱动程序版本：550.54.14
• CUDA版本：12.4
• cuDNN版本：8.9.7
• 操作系统：Debian 11
• Python版本：3.9.2
• PyTorch版本：2.2.1+cu121

问题表现

1. 内存使用量随时间持续增长，最终耗尽系统64GB内存
2. 模型转换过程极其缓慢，即使简化后也需要约10分钟
3. 转换过程中显存占用高达1.5GB
4. 原始模型存在类型不匹配问题（int64和int32），修复后问题依然存在

技术分析

经过深入分析，这个问题可能由以下几个因素导致：

1. 循环结构处理：模型中的循环结构可能导致TensorRT优化器进入无限优化循环
2. 内存管理缺陷：TensorRT 8.6.1版本在构建引擎时可能存在内存释放不及时的问题
3. 张量拼接操作：频繁的torch.cat操作可能触发TensorRT内部的内存分配异常

解决方案

NVIDIA官方已确认此问题并提交了内部bug报告（编号4601836）。根据官方回复，该问题已在最新小版本更新中得到修复。建议用户采取以下措施：

1. 升级到TensorRT的最新稳定版本
2. 对于必须使用8.6.1版本的用户，可以尝试以下临时解决方案：
   • 限制模型输入尺寸
   • 减少循环次数
   • 增加内存监控，在内存达到阈值时重启进程

最佳实践建议

1. 在模型转换前，先使用ONNX Runtime验证模型有效性
2. 对于包含循环结构的模型，建议先在小规模输入上测试转换性能
3. 监控转换过程中的内存使用情况，设置合理的超时机制
4. 考虑将复杂循环结构拆分为多个子图，降低转换复杂度

总结

TensorRT作为高性能推理引擎，在大多数情况下表现优异，但在处理特定模型结构时仍可能出现问题。开发者在使用过程中应当注意监控系统资源使用情况，并及时更新到稳定版本以获得最佳体验和性能。对于类似的内存泄漏问题，保持与官方社区的沟通并及时反馈是解决问题的有效途径。