在 Linux 系统上训练的时候，在训练中发现程序被 killed 了。
查看 killed 的原因：
dmesg | egrep -i -B100 'killed process'
会得到如下的信息
可以看到这个程序被 killed 的原因是 out of memory， total-vm 指的是虚拟内存，anon-rss 是占用的物理内存。
接下来查找代码中可能导致内存爆满的操作。检查思路有：
- 不要在循环中累加。比如计算 loss， pytorch 里面的 loss 直接相加会让梯度累加，导致内存保证。可以用
total_loss+=loss.item()或者total_loss+=float(loss) - 不要保存不需要的东西。局部变量以及对象的成员变量会在运行到变量范围之外才会解除分配，用不到的东西就删掉，减少内存占用。
- 不要用太大的线形层。一个线性层
nn.Linear(m,n)的内存是，在训练的时候需要的内存至少是 2 倍的 mn，因为还需要存梯度。分类任务的最后一层线性层可以用 global average pooling 来代替，参数更少。
参考：
https://pytorch.org/docs/stable/notes/faq.html
https://blog.csdn.net/m0_37798080/article/details/119870459
网友评论