之前使用torch.cuda.empty_cache()都不会出现问题,当修改了使用的显卡号就出现问题了。
在第一轮训练nvidia-smi看显存是够的,但是测试结束后开始第二轮训练就报错了。报错如下:
File "train.py", line 64, in train
torch.cuda.empty_cache()
File "torch/cuda/memory.py", line 114, in empty_cache
torch._C._cuda_emptyCache()
RuntimeError: CUDA error: out of memory
之前的经验一般是batch_size设置太大了,显存不够,导致训练的时候爆显存,而这次是因为释放显存而爆显存?很奇怪的操作。。。
后来才发现,在多卡训练的时候容易出现这个问题,单卡不会,因为torch.cuda.empty_cache()默认是给GPU0释放显存,而此时GPU0跑满了,无显存可以释放,就会导致显存溢出。
解决方法为:
with torch.cuda.device(device):
torch.cuda.empty_cache()
device即你设置的显卡号,这样就不会出错了