torch.cuda.empty_cache()导致RuntimeError: CUDA error: out of memory.

用双卡训练时,当GPU0接近跑满时,在GPU1中如果有

torch.cuda.empty_cache()

会导致:

torch._C._cuda_emptyCache()
RuntimeError: CUDA error: out of memory.

出现这个问题是当使用GPU1训练时,torch.cuda.empty_cache()默认是给GPU0释放显存,而此时GPU0接近跑满,且无显存可释放,直接使用这句话就会导致显存溢出。解决方法为:

    with torch.cuda.device('cuda:1'):
        torch.cuda.empty_cache()

即指定给GPU1回收显存,这样就可以正常运行了。

参考资料:
1.https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898
2.https://github.com/PyTorchLightning/pytorch-lightning/issues/3016

你可能感兴趣的:(深度学习,pytorch,深度学习,神经网络,python)