如何释放Pytorch CUDA中的GPU内存

Question

我正在使用Colab和Pytorch CUDA进行我的深度学习项目，并面临无法释放GPU的问题。我在这里读过一些相关的帖子，但他们没有解决我的问题。请指导我如何释放GPU内存。提前谢谢你

Answer 1

试试这个:

torch.cuda.empty_cache()

或:

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

Answer 2

这段代码可以做到这一点。GPU的索引为[0,1，…]所以如果你只有一个，那么gpu_index为0。注意，这通过杀死底层tensorflow会话来清除GPU。你的模型，数据等从GPU中清除，但它不会重置Colab/Jupyter会话上的内核。因为它会清除会话，所以你不能在运行期间使用这个来清除内存。

from numba import cuda 
def clear_GPU(gpu_index):
    cuda.select_device(gpu_index) 
    cuda.close()

Install numba ("pip Install numba")…上次我尝试conda给了我问题，所以使用pip。这是一个方便的b/c numba开发人员已经采取了适当执行一些低级CUDA方法的麻烦，所以我想你可以做同样的，如果你有时间。