如何解决Python中PyTorch显存泄露问题_通过torch.cuda.empty_cache释放

torch.cuda.empty_cache()仅释放CUDA缓存中空闲块，不释放被张量占用的显存；显存泄漏需通过控制张量生命周期、切断计算图引用和及时del+gc.collect()解决。

torch.cuda.empty_cache() 不是万能解药，它只清空当前设备上 CUDA 缓存分配器里的空闲块，对真正泄漏的显存无效，反而可能拖慢高频推理。

torch.cuda.empty_cache() 为什么经常没用

• 它不释放正在被张量占用的显存（即 memory_allocated() 部分），只动 memory_reserved() 中已释放但未归还给驱动的缓存块
• 如果你没主动删掉中间张量、没断开计算图引用，empty_cache() 调多少次都看不到显存回落
• 在多卡场景下，它默认只作用于 current_device，若你用 torch.cuda.set_device(1) 切过卡，却在主卡上调用，等于白清

常见错误现象：

• 每次 model(x) 后立刻跟 torch.cuda.empty_cache()，nvidia-smi 数值纹丝不动
• 多卡训练时只在 rank 0 清缓存，其他卡显存持续上涨

正确做法：

Python 3.14.3

微软官方的 Python 扩展，是 VS Code 安装量最高的扩展（209M+）。集成 IntelliSense（通过 Pylance）、调试（通过 Python Debugger）、代码检查、格式化、重构和单元测试等功能。支持 Jupyter Notebook、虚拟环境管理和多 Python 版本切换。

• 先确认是不是真泄漏：插两行监控

  print(f"Allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB")
  print(f"Reserved: {torch.cuda.memory_reserved()/1e9:.2f} GB")

• 若 allocated 持续涨 → 是代码级泄漏，不是缓存问题
• 若 reserved 明显 > allocated → 缓存堆积，empty_cache() 才有意义

推理循环中必须做的三件事（比 empty_cache() 更关键）

• 用 with torch.no_grad(): 包住前向过程，否则 requires_grad=True 会偷偷建图、留引用
• 所有中间张量（尤其是 loss、指标、临时输出）必须显式脱离图：
  • 用 .item() 取标量值（不要 += loss）
  • 用 .detach().cpu() 搬到 CPU 或转 numpy（别直接 .numpy()，会报错）
  • 输出 tensor 不再需要时，立刻 del pred，再调 gc.collect()
• DataLoader 的 num_workers > 0 时，避免传入 list/dict 等非 tensor 类型，否则子进程会 copy-on-access 持有原始内存

什么时候该调 torch.cuda.empty_cache()

• 在 epoch 结束后 或 验证/测试阶段切换前后，且已确保所有中间张量被 del + gc.collect()
• 单卡调试时，配合 torch.cuda.memory_summary() 看缓存块分布，确认清的是“空闲但未归还”的那部分
• 不要放在 for x, y in dataloader: 循环体内，高频调用会让小张量分配变慢（缓存池清空后得重新向 driver 申请页）

empty_cache() 是个“事后清理员”，不是“泄漏修复器”。真正管用的永远是：控制张量生命周期、切断计算图引用、让 Python 垃圾回收能真正触达对象。显存爬升时先盯 memory_allocated()，别被 nvidia-smi 的总数骗了。