Python程序PyTorch显存泄漏怎么办_利用torch.cuda.empty_cache清理

torch.cuda.empty_cache() 仅释放未被张量引用的缓存显存，不回收仍被变量或模型持有的显存；需配合 del、zero_grad() 和 no_grad() 才能有效释放。

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

简单来说，这个函数的作用范围很有限。它只负责清理CUDA缓存分配器里那些“无主”的显存块，本质上不是一个垃圾回收器。换句话说，只要你的torch.Tensor还被某个Python变量引用着，或者模型参数还被nn.Module结构体包裹着，empty_cache()就对它们束手无策。

一个典型的场景就是：运行nvidia-smi时发现显存占用一路攀升，反复调用torch.cuda.empty_cache()却收效甚微，数字几乎不动。问题出在哪？

• 检查中间变量：比如在训练循环里，如果每次前向传播都产生一个新的output张量，但没有显式地del它或用新值覆盖，这些张量就会一直“活着”。
• 验证梯度状态：推理时如果没用with torch.no_grad():把代码包起来，PyTorch就会默默构建计算图，这些中间结果同样会占用显存。
• 分清模型模式与内存管理：别忘了，model.eval()只是关闭了Dropout和BatchNorm的训练模式，它可不会帮你释放已经分配好的参数和缓存。

真正能释放显存的三步操作

指望单靠empty_cache()解决问题，无异于扬汤止沸。真正有效的方法，需要配合Python的引用计数和PyTorch的运行时管理，三步走：

• 手动切断引用：对于不再需要的张量，比如推理后的输出output、计算完的loss，直接使用del output, loss。尤其在多轮推理的场景下，别把希望全寄托在解释器的自动回收上。
• 清空计算图：在loss.backward()之后，立即调用optimizer.zero_grad()。否则，梯度（grad）会一直引用着整个计算图，导致相关显存无法释放。
• 最后调用缓存清理：将torch.cuda.empty_cache()放在所有del和zero_grad()操作之后。需要注意的是，建议只在调试或批处理任务的间歇期使用它，不要把它塞进每一步的前向传播里。

来看一个更清晰的示例片段：

立即学习“Python免费学习笔记（深入）”；

for x, y in dataloader:
    with torch.no_grad():
        pred = model(x)
        # ... 计算指标
    del pred  # 关键：主动删掉输出张量
    torch.cuda.empty_cache()  # 放在这里才可能生效

哪些情况 empty_cache() 反而有害

这个函数并非无害的“万能药”。它会强制清空CUDA缓存分配器中的空闲内存块，导致后续需要分配新张量时，不得不重新向驱动申请内存页。如果调用频率过高，不仅释放不了多少显存，反而会显著拖慢程序速度。

• 避免高频调用：千万不要在每个forward()后面都加一句empty_cache()。
• 注意多卡环境：它只对当前设备（current_device）生效。如果你用torch.cuda.set_device()切换过GPU，要确保清理的是正确的目标卡。
• 使用更精准的工具监控：与其依赖nvidia-smi的粗略数据，不如搭配torch.cuda.memory_summary()。运行一次后，重点关注allocated（已分配）和reserved（预保留）之间的差值，这更能反映PyTorch内部的真实使用情况。

排查泄漏的最小可行路径

遇到显存问题，先别急着大改代码。按照一个系统化的路径来排查，往往事半功倍。

• 加装监控点：在训练循环的关键位置（如每个epoch开始/结束时）插入内存查询：print(torch.cuda.memory_allocated()/1024**3)，观察显存占用的增长趋势。
• 调整分配策略测试：可以尝试设置环境变量os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'。这能限制缓存分配器持有过大的内存块，有助于识别是否因内存碎片导致“隐形”占用。
• 隔离模型问题：如果怀疑是模型本身有泄漏，可以尝试将模型移到CPU：model.cpu()，然后删除del model，再调用empty_cache()，观察显存是否如预期回落。

最后，有一个极易被忽略的“坑”：DataLoader设置pin_memory=True并结合GPU张量预加载时，会在子进程（worker）中提前占用显存。这部分显存不受主进程的empty_cache()管理，需要特别注意。