最近用Pytorch跑模型,在自己电脑上跑的挺好的,但是一送到服务器上去跑就挂了。一开始以为是batchsize开太大了爆炸了,后面调小了batchsize依旧爆炸,才意识到可能不是这里的问题,应该是出现内存泄漏了。
于是我边跑模型边盯着服务器的RAM占用,果不其然,在模型开始跑了以后,RAM占用还是在不断上升,直到完全占满了内存和swap,产生killed。
知道是发生了内存泄漏,我还是比较惊奇的。这可是Python啊,又不是C,一不小心指针就得出点什么问题。后来经过在网上搜寻,我找到了罪魁祸首,看如下代码:
d_loss_all += d_loss
g_loss_all += g_loss
没错,就是这里出了问题,根据网上的说法,这里的loss是包含自动求导的,如此直接相加会导致autograd历史不断被保存,内存未被释放,最终造成内存爆炸。
解决起来很简单,我们需要一个普通的float数相加就可以了,将代码改为:
d_loss_all += d_loss.item()
g_loss_all += g_loss.item()
完美解决我的问题。
最后,将swap中的内存转到真正的内存上来
swapoff -a && swapon -a
-=||=-收藏