如何在PyTorch中清理CUDA显存垃圾，解决显存占用不释放问题？

你提到的Python垃圾回收器无法作用于CUDA设备是对的：Python gc仅负责回收CPU内存中的Python对象引用，CUDA显存由PyTorch的CUDA runtime单独管理，只有当所有指向某块CUDA显存的Python侧引用全部被回收，且PyTorch主动将该块显存标记为可释放时，torch.cuda.empty_cache()才会将这部分显存还给CUDA设备。如果出现清理后显存仍被占用的情况，按照以下步骤排查处理即可：

1. 排查遗漏的CUDA张量引用 #

很多时候显存无法释放是因为存在未被注意到的张量引用，没有被gc回收：

• 张量可能被隐藏在全局列表、字典、闭包、装饰器缓存、模型的注册钩子（如register_forward_hook、register_backward_hook）中，这些引用不会因为删除了表面的model、input变量就被回收
• 可以执行以下代码遍历所有现存的CUDA张量，定位残留引用：

import torch
import gc

for obj in gc.get_objects():
    try:
        if torch.is_tensor(obj) and obj.is_cuda:
            print(f"张量类型：{type(obj)}, 尺寸：{obj.size()}, 设备：{obj.device}")
    except Exception:
        pass

• 定位到残留张量后，手动删除对应的持有对象，再依次执行gc.collect()、torch.cuda.empty_cache()即可释放对应显存

2. 重置CUDA上下文 #

如果确认所有Python侧的CUDA张量引用都已经删除，显存仍然被占用，说明CUDA上下文本身存在残留统计或持有：

• 先执行torch.cuda.reset_peak_memory_stats()重置显存统计数据，再调用torch.cuda.empty_cache()尝试释放
• 如果以上操作无效，Colab环境下可以执行以下命令强制杀掉所有关联CUDA的Python进程，彻底释放显存：
  !nvidia-smi | grep -i python | awk '{print $5}' | xargs kill -9
  注意该命令会清空当前运行时的所有变量，需要重新执行代码。

3. 避免显存残留的前置规范 #

日常使用中遵循以下规则可以大幅降低显存残留的概率：

• 尽量不要在全局作用域定义大的CUDA张量，用完的张量及时用del删除
• 推理代码全程用with torch.no_grad()包裹，避免生成不必要的计算图和梯度缓存
• 循环执行多轮任务时，尽量复用张量空间，不要重复创建新的大尺寸CUDA张量
• 最彻底的释放方式是重启Python运行时，Colab中可直接点击「代码执行程序」-「重新启动代码执行程序」，所有显存占用会被清空。

内容的提问来源于stack exchange，提问作者Лолаев Георгий