PyTorch CUDA内存溢出｜RuntimeError内存不足问题排查与解决

典型报错 #

RuntimeError: CUDA out of memory. Tried to allocate 916.00 MiB (GPU 0; 6.00 GiB total capacity; 4.47 GiB already allocated; 186.44 MiB free; 4.47 GiB reserved in total by PyTorch; 6.00 GiB total capacity; 4.47 GiB already allocated; 186.44 MiB already allocated. 186.44 MiB free; 4.47 GiB reserved in total by PyTorch)

分步排查步骤 #

• 检查GPU进程占用：执行nvidia-smi命令，确认是否有未关闭的PyTorch进程、其他深度学习框架进程或GPU应用抢占显存，残留进程是常见的隐性占用原因。
• 排查张量的不必要保留：
  • 确认是否在计算后未及时释放无用张量，比如未调用del tensor_name手动释放，或未触发Python垃圾回收。
  • 检查推理/测试阶段是否未禁用自动梯度计算，autograd保留的计算图会占用大量显存。
  • 排查循环中是否重复创建新张量但未清理，或inplace操作不当导致的张量残留。
• 检查数据加载配置：即使模型极小，过大的batch_size会导致输入数据占用大量显存；同时确认数据预处理过程中是否在GPU上创建了不必要的大张量。
• 验证环境兼容性：检查PyTorch版本与CUDA驱动、CUDA Toolkit版本是否匹配，版本不兼容可能引发显存分配异常或泄漏。
• 查看详细内存统计：执行torch.cuda.memory_summary(device=0)获取精确的显存使用明细，区分已分配、缓存、保留显存的占比，定位异常占用项。

极小模型仍出现问题的解决方法 #

• 强制清理缓存显存：执行torch.cuda.empty_cache()，释放PyTorch已申请但未使用的缓存显存（无法释放正在使用的显存）。
• 禁用自动梯度计算：在推理/测试代码块外包裹with torch.no_grad():，避免计算图占用显存；训练阶段可对不需要梯度的张量调用tensor.detach()截断梯度跟踪。
• 清理隐性内存占用：检查全局变量是否持有未使用的大张量并及时删除；若使用Jupyter Notebook，直接重启内核清理残留进程的显存占用。
• 限制显存分配比例：通过torch.cuda.set_per_process_memory_fraction(0.7, device=0)设置当前进程可使用的显存比例（示例为70%），避免进程占满整个GPU显存。
• 排查内存泄漏：若问题出现在循环训练/推理中，逐步注释代码，定位是否存在每次循环都累积占用显存的操作，比如未在循环内释放临时张量。

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