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

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

2026-05-28日常编程116547

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 的总数骗了。