使用模型后如何清除GPU内存？

Question

我正在尝试在使用完模型后释放 GPU 内存。

• 我在创建和训练模型之前检查了 nvidia-smi：402MiB / 7973MiB
• 创建和训练模型后，我再次使用 nvidia-smi 检查 GPU 内存状态：7801MiB / 7973MiB
• 现在我尝试通过以下方式释放 GPU 内存：

del model
torch.cuda.empty_cache() 
gc.collect()

并再次检查 GPU 内存：2361MiB / 7973MiB

• 如您所见，并未释放所有 GPU 内存（我预计会得到 400~MiB / 7973MiB）。
• 我只能通过终端释放 GPU 内存（sudo fuser -v /dev/nvidia* 和 kill pid）

有没有办法在我使用完模型后释放 GPU 内存？

Answer 1

发生这种情况是因为 pytorch 为快速内存分配保留了 gpu 内存。要了解更多信息，请参阅 pytorch memory management。要解决此问题，您可以使用以下代码：

from numba import cuda
cuda.select_device(your_gpu_id)
cuda.close()

但是，这有一个问题。它完全关闭了 GPU。所以，你不能在不重新启动一切的情况下开始训练。