梯度累积优化—切实可行的OOM解决方案

首先关于OOM，之前写过一篇文章：https://www.jianshu.com/p/de56e12a744e

常见的GPU的OOM的解决方案

1、减小batchsize

2、减小模型

3、增大显卡

4、并行GPU

常见的GPU的OOM的解决方案的实施缺陷

首先方法1，理论上说在合适的学习率下，batchsize=1也能有比较好的学习效果，但现实多残酷啊，如何找到合适的学习率？在实际的训练中效果真会好吗？现实是batchsize非常小的时候（极限情况下=1，也就说每次用一个样本的训练来更新梯度），此时方差很大，导致我们的模型训练起来非常慢（loss前期是在下降的，但是下降速度非常慢，后期有可能不降反升），而且实验证明了，batchsize小的效果真不如batchsize大的效果，不论在loss的下降速度还是准确率上。

方法2呢，我们大多时候是不想要减小模型的，有时候我们想复现他人的论文模型，模型减小就失去了对比的意义。

方法3，增大显卡，当然这是由老板决定的。

方法4，并行GPU是可行的，如果你有多个GPU的话。

干货：新的OOM解决方案——梯度累积优化

当训练模型时出现OOM，我们显存有限，此时又不想减小网络，我们只能把batchsize降低一些，一降再降，从128降到64，32，16。实验证明，batchsize=128和batchsize=16，前者准确率可能比后者高几个千分点，甚至几个百分点。而此时，梯度累积优化方法就派上用场了，他可以以时间换空间，相当于增大了batchsize。

所谓梯度累积，是说我们更新参数所用的梯度，实际上是多个样本算出来的梯度的平均值，以batch_size=128为例，可以一次性算出128个样本的梯度然后平均，也可以每次算16个样本的平均梯度，然后缓存累加起来，算够了8次之后，然后把总梯度除以8，然后才执行参数更新。当然，必须累积到了8次之后，用8次的平均梯度才去更新参数，不能每算16个就去更新一次，不然就是batch_size=16了。这就是梯度累积优化的含义所在。

当然了，这种方法并不能加快模型的训练时间，但是却可以尽情的增大batchsize而不造成OOM。当然，如果你有多个GUP，可以做并行计算，那就可以把任务分配到多个GPU上并行计算，最后再累加计算平均值，时间和效率都有了。

这个优化器的修改，使得小batch size能起到大batch size的效果，前提是模型不包含Batch Normalization，因为Batch Normalization在梯度下降的时候必须用整个batch的均值方差。所以如果你的网络用到了Batch Normalization，想要准确达到大batch size的效果，目前唯一的方法就是加显存/加显卡。

梯度累积优化实现

那么问题来了，这种优化器怎么实现呢？基于Keras，苏大神开源了梯度累积优化器的实现代码，方便好用。比如说128分成了16*8，batchsize=16，算8次更新一次梯度，那就是前7次更新都为0，第8次才更新。

实现代码在这里：https://github.com/bojone/accum_optimizer_for_keras

使用代码也很简单：

梯度累积优化器使用代码

以上例子等价于直接使用batch_size=100的Adam优化器（代价就是你跑了10个epoch，实际上只相当于batch_size=100跑了1个epoch）。

以上！终于找到一种比较合理的OOM解决方案，又是开心的一天鸭！