TensorRT 在 RTX2000 GPU 上运行 ONNX 模型的内存优化策略

问题背景

在使用 TensorRT 8510 版本将 ONNX 模型转换为 TensorRT 引擎时，特别是在 RTX2000 GPU 上运行时，开发者可能会遇到内存不足的警告和错误。这些错误信息表明 TensorRT 在尝试分配大块 GPU 内存时遇到了困难，导致某些优化策略被跳过。

错误现象分析

当执行 polygraphy 工具进行 ONNX 到 TensorRT 的转换时，系统会输出以下关键警告和错误信息：

1. INT64 权重类型警告：ONNX 模型使用了 INT64 类型的权重，而 TensorRT 原生不支持 INT64，系统会自动尝试将其降级为 INT32 类型。

2. 内存分配错误：

   • 系统尝试分配 8GB (8589934592 字节) 的 GPU 内存失败
   • 由于内存不足，某些优化策略(tactic)被跳过
   • 建议通过 IBuilderConfig::setMemoryPoolLimit() 减少工作空间大小

3. 构建结果：尽管有内存警告，引擎构建最终在约126秒后完成

技术原理

TensorRT 在构建引擎时会尝试多种优化策略(tactics)来提高推理性能。每种策略可能需要不同大小的 GPU 内存作为工作空间(workspace)。当可用内存不足时：

1. TensorRT 会跳过需要更多内存的优化策略
2. 如果内存严重不足，整个进程可能会被系统终止
3. 内存需求取决于模型复杂度、批次大小和优化级别

解决方案与优化建议

1. 调整工作空间内存限制

通过 IBuilderConfig::setMemoryPoolLimit() 显式设置更小的工作空间内存限制：

config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 设置为1GB

2. 优化模型输入尺寸

减小模型的最大输入尺寸可以显著降低内存需求：

profile = builder.create_optimization_profile()
profile.set_shape("input_name", min=(1,3,224,224), opt=(4,3,224,224), max=(8,3,224,224))
config.add_optimization_profile(profile)

3. 使用更低精度的数据类型

考虑使用 FP16 或 INT8 量化来减少内存占用：

config.set_flag(trt.BuilderFlag.FP16)
# 或
config.set_flag(trt.BuilderFlag.INT8)

4. 分批处理策略

对于大模型，可以考虑：

1. 减小推理时的批次大小
2. 实现流式处理，避免一次性加载整个模型

RTX2000 GPU 特别注意事项

RTX2000 系列 GPU 的内存相对有限，使用时需要特别注意：

1. 监控 GPU 内存使用情况：使用 nvidia-smi 工具
2. 关闭不必要的进程释放 GPU 内存
3. 考虑使用更轻量级的模型架构

结论

在资源受限的 GPU 如 RTX2000 上使用 TensorRT 时，内存管理是关键。通过合理配置工作空间大小、优化模型参数和使用适当精度，可以成功构建和运行 TensorRT 引擎，即使遇到初始的内存分配警告。理解这些警告的含义有助于开发者做出正确的优化决策，平衡性能和资源使用。