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

Question

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

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

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

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

• 正如您所看到的，并非所有 GPU 内存都被释放（我预计会获得400~MiB / 7973MiB）。
• 我只能通过终端（sudo fuser -v /dev/nvidia* 和killpid）释放GPU内存

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

Answer 1

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

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

然而，这有一个问题。它完全关闭 GPU。因此，如果不重新启动一切，就无法开始训练。