Pytorch在训练模型时,显存不断增大,导致CUDA out of memory。

一般遇到CUDA out of memory错误,把batch_size设小即可解决。但有时会遇到,无论如何调小batch_size,都还是会出现CUDA out of memory的问题。

问题得具体情况:一开始模型能跑,跑了一会,出现CUDA out of memory错误,这种情况是因为显存不断增大导致的。

解决办法:
在保存打印数据的时候,要把数据转为cpu的格式储存。有时候我们会设置每隔5或10个epoch打印一次loss,下面的代码有时候会导致CUDA out of memory错误:

total_loss += loss

因为loss还是在显存中,这样累加的语句会导致显存不断加大,最终出现CUDA out of memory的错误,改为如下语句即可解决(该语句把loss改为cpu储存):

total_loss += loss.cpu().detach().numpy().item()

另一种情况(开头介绍的),当程序一开时运行的时候就出现CUDA out of memory错误,这种情况一般是batch_size太大,超过GPU显存导致的,此时改小batch_size就可以。

深入探究请参考:https://zhuanlan.zhihu.com/p/85838274

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