如何解决CUDA Out of Memory错误并优化CUDA11.3推理速度？

问题背景 #

• 运行检测+识别模型推理，CUDA10.2环境下耗时15分钟，但硬件不支持CUDA10.2，切换到CUDA11.3后耗时长达3小时，且出现OOM错误：

RuntimeError: CUDA out of memory. Tried to allocate 2.05 GiB (GPU 0; 5.81 GiB total capacity; 2.36 GiB already allocated; 1.61 GiB free; 2.38 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

• 当前环境：
  • cudatoolkit 11.3.1 h2bc3f7f_2
  • pytorch 1.10.0 py3.7_cuda11.3_cudnn8.2.0_0 pytorch
  • torchvision 0.11.0 py37_cu113 pytorch
• 已尝试调整检测/识别的batch_size，问题未解决

一、先解决CUDA Out of Memory问题 #

OOM会迫使你使用极小batch甚至单样本推理，是速度慢的核心原因之一，试试这些方案：

1. 优化显存分配策略：
   启动CLI前设置环境变量，减少显存碎片：
   • Linux/macOS：export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
   • Windows：set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
2. 半精度推理：
   如果CLI支持，添加--fp16或--half参数；如果是自定义脚本，用torch.cuda.amp.autocast()包裹推理逻辑，同时把模型转成torch.float16()，显存占用直接减半，对5.8G显存的卡极为有效。
3. 禁用梯度计算：
   推理时必须确保关闭梯度，自定义脚本用torch.no_grad()上下文管理器；如果是第三方CLI，检查是否有--no-grad类参数，避免显存被梯度占用。
4. 定期清理显存缓存：
   在推理循环间隙（比如每处理完一批数据）调用torch.cuda.empty_cache()，释放无用显存。
5. 极端情况：拆分任务：
   把待处理的数据集分成更小的子集，逐子集推理，每处理完一个子集就清一次缓存。

二、优化CUDA 11.3下的推理速度 #

解决OOM后，再针对性提升速度：

1. 升级PyTorch版本：
   你当前用的PyTorch1.10.0对CUDA11.3的优化有限，建议升级到1.12.x或1.13.x（均兼容CUDA11.3），新版本会同步升级cudnn，对CUDA11.x的硬件适配更好。
2. 启用cudnn基准模式：
   在推理前设置torch.backends.cudnn.benchmark = True，让cudnn自动选择最优卷积算法，固定输入尺寸的场景下能显著提速。
3. 消除CPU-GPU传输瓶颈：
   确保模型、输入数据提前全部移至GPU，不要在推理循环内反复执行.to('cuda')，减少数据传输耗时。
4. TensorRT模型优化：
   将PyTorch模型转换为TensorRT引擎，针对你的GPU做硬件级优化，TensorRT对CUDA11.x的支持完善，能大幅提升CNN模型的推理速度。如果CLI不支持导出，可以自行编写脚本完成转换。
5. 优化预处理/后处理：
   用nvidia-smi查看推理时的GPU利用率，如果利用率偏低，说明CPU预处理拖慢了速度：
   • 将图片缩放、归一化等预处理操作移至GPU执行
   • 使用多线程/多进程做CPU预处理，避免GPU等待

内容的提问来源于stack exchange，提问作者aarya