如何在不重启内核的情况下在 PyTorch 模型训练后清除 GPU 内存

Question

我正在 Jupyter-Lab 笔记本上训练 PyTorch 深度学习模型，在 Tesla K80 GPU 上使用 CUDA 进行训练。在进行训练迭代时，会使用 12 GB 的 GPU 内存。我通过保存模型检查点来完成训练，但想继续使用笔记本进行进一步分析（分析中间结果等）。

但是，这些 12 GBnvtop在完成训练后继续被占用（如从中看到的）。我想释放这个内存，以便我可以将它用于其他笔记本。

到目前为止，我的解决方案是重新启动此笔记本的内核，但这并不能解决我的问题，因为到目前为止我无法继续使用相同的笔记本及其各自计算的输出。

Answer 1

with torch.no_grad():
    torch.cuda.empty_cache()

Answer 2

到目前为止，Cuda 方面的答案是正确的，但 ipython 方面也存在问题。

当您在笔记本环境中出现错误时，ipython shell 会存储异常的回溯，以便您可以使用%debug. 问题是这需要将导致错误的所有变量保存在内存中，并且它们不会被gc.collect(). 基本上你所有的变量都卡住了，内存泄漏了。

通常，引发新异常会释放旧异常的状态。所以尝试类似的方法1/0可能会有所帮助。但是，使用 Cuda 变量可能会变得很奇怪，有时在不重新启动内核的情况下无法清除 GPU 内存。

有关更多详细信息，请参阅这些参考资料：

Answer 3

如果您只是将使用大量内存的对象设置为None这样：

obj = None

之后你打电话

gc.collect() # Python thing

这是您可以避免重新启动笔记本的方法。

如果您仍然希望从 Nvidea smi 或 nvtop 中清楚地看到它，您可以运行：

torch.cuda.empty_cache() # PyTorch thing

清空 PyTorch 缓存。

是的，“torch.cuda.empty_cache()”可能会帮助您清楚地看到它。 (3认同)
我尝试了 `model = None` 和 `gc.collect()` 但它没有清除任何 GPU 内存 (2认同)
它以同样的顺序对我有用。1.- model = None, 2.- gc.collect(), 3.- torch.cuda.empty_cache() (2认同)

Answer 4

我自己从未使用过 PyTorch，但 Google 有几个结果，基本上都是一样的。torch.cuda.empty_cache()

`torch.cuda.empty_cache()` 清除了大部分已用内存，但我仍然有 2.7GB 正在使用。可能是模型占用了内存，但我不知道它有多清晰。我尝试了其他答案中的“model = None”和“gc.collect()”，但它不起作用。 (9认同)