Github地址：https://github.com/bojone/accum_optimizer_for_keras
扯淡
在一两年之前，做 NLP 任务都不用怎么担心 OOM 问题，因为相比 CV 领域的模型，其实大多数 NLP 模型都是很浅的，极少会显存不足。幸运或者不幸的是，Bert 出世了，然后火了。Bert 及其后来者们（GPT-2、XLNET 等）都是以足够庞大的 Transformer 模型为基础，通过足够多的语料预训练模型，然后通过 fine tune 的方式来完成特定的 NLP 任务。
即使你很不想用 Bert，但现在的实际情况是：你精心设计的复杂的模型，效果可能还不如简单地 fine tune 一下 Bert 好。所以不管怎样，为了跟上时代，总得需要学习一下 Bert 的 fine tune。
问题是“不学不知道，一学吓一跳”，只要任务稍微复杂一点，或者句子长度稍微长一点，显存就不够用了，batch size 急剧下降——32？16？8？一跌再跌都是有可能的。
这不难理解，Transformer 基于 Attention，而 Attention 理论上空间和时间复杂度都是
，虽然在算力足够强的时候，Attention 由于其并行性还是可以表现得足够快，但是显存占用量是省不了了，
意味着当你句子长度变成原来的 2 倍时，显存占用基本上就需要原来的 4 倍，这个增长比例肯定就容易 OOM 了。
而更不幸的消息是，大家都在 fine tune 预训练 Bert 的情况下，你 batch_size=8 可能比别人 batch_size=80 低好几个千分点甚至是几个百分点，显然这对于要刷榜的读者是很难受的。难道除了加显卡就没有别的办法了吗？
正事
有！通过梯度缓存和累积的方式，用时间来换取空间，最终训练效果等效于更大的 batch size。因此，只要你跑得起 batch_size=1，只要你愿意花 n 倍的时间，就可以跑出 n 倍的 batch size 了。
梯度累积的思路，在之前的文章“让Keras更酷一些！”：小众的自定义优化器已经介绍了，当时称之为“软 batch（soft batch）”，本文还是沿着主流的叫法称之为“梯度累积（accumulate gradients）”好了。
所谓梯度累积，其实很简单，我们梯度下降所用的梯度，实际上是多个样本算出来的梯度的平均值，以 batch_size=128 为例，你可以一次性算出 128 个样本的梯度然后平均，我也可以每次算 16 个样本的平均梯度，然后缓存累加起来，算够了 8 次之后，然后把总梯度除以 8，然后才执行参数更新。当然，必须累积到了 8 次之后，用 8 次的平均梯度才去更新参数，不能每算 16 个就去更新一次，不然就是 batch_size=16 了。
刚才说了，在之前的文章的那个写法是有误的，因为用到了：
K.switch(cond, K.update(p, new_p), p)
来控制更新，但事实上这个写法不能控制更新，因为 K.switch 只保证结果的选择性，不保证执行的选择性，事实上它等价于：
cond * K.update(p, new_p) + (1 - cond) * p也就是说不管 cond 如何，两个分支都是被执行了。事实上 Keras 或 Tensorflow“几乎”不存在只执行一个分支的条件写法（说“几乎”是因为在一些比较苛刻的条件下可以做到），所以此路不通。
不能这样写的话，那只能在“更新量”上面下功夫，如前面所言，每次算 16 个样本的梯度，每次都更新参数，只不过 8 次中有 7 次的更新量是 0，而只有 1 次是真正的梯度下降更新。
很幸运的是，这种写法还可以无缝地接入到现有的 Keras 优化器中，使得我们不需要重写优化器！详细写法请看：
https://github.com/bojone/accum_optimizer_for_keras
具体的写法无外乎就是一些移花接木的编程技巧，真正有技术含量的部分不多。关于写法本身不再细讲，如果有疑问欢迎讨论区讨论。
注：这个优化器的修改，使得小 batch size 能起到大 batch size 的效果，前提是模型不包含 Batch Normalization，因为 Batch Normalization 在梯度下降的时候必须用整个 batch 的均值方差。所以如果你的网络用到了 Batch Normalization，想要准确达到大 batch size 的效果，目前唯一的方法就是加显存/加显卡。
实验
至于用法则很简单：
opt = AccumOptimizer(Adam(), 10) # 10是累积步数
model.compile(loss='mse', optimizer=opt)
model.fit(x_train, y_train, epochs=10, batch_size=10)这样一来就等价于 batch_size=100 的 Adam 优化器了，代价就是你跑了 10 个 epoch，实际上只相当于 batch_size=100 跑了 1 个 epoch，好处是你只需要用到 batch_size=10 的显存量。
可能读者想问的一个问题是：你怎么证明你的写法生效了？也就是说你怎么证明你的结果确实是 batch_size=100 而不是 batch_size=10？
为此，我做了个比较极端的实验，代码在这里：
https://github.com/bojone/accum_optimizer_for_keras/blob/master/mnist_mlp_example.py
代码很简单，就是用多层 MLP 做 MNIST 分类，用 Adam 优化器， fit 的时候 batch_size=1。优化器有两个选择，第一个是直接 Adam() ，第二个是 AccumOptimizer(Adam(), 100) ：
如果是直接 Adam() ，那 loss 一直在 0.4 上下徘徊，后面 loss 越来越大了（训练集都这样），val 的准确率也没超过 97%；
如果是 AccumOptimizer(Adam(), 100) ，那么训练集的 loss 越来越低，最终降到 0.02 左右，val 的最高准确率有 98%+；
最后我比较了直接 Adam() 但是 batch_size=100 的结果，发现跟 AccumOptimizer(Adam(), 100) 但是 batch_size=1 时表现差不多。
这个结果足以表明写法生效了，达到了预期的目的。如果这还不够说服力，我再提供一个训练结果作为参考：
在某个 Bert 的 fine tune 实验中，直接用 Adam() 加 batch_size=12，我跑到了 70.33% 的准确率；我用 AccumOptimizer(Adam(), 10) 加 batch_size=12（预期等效 batch size 是 120），我跑到了 71% 的准确率，提高了 0.7%，如果你在刷榜，那么这 0.7% 可能是决定性的。
结论
终于把梯度累积（软 batch）正式地实现了，以后用 Bert 的时候，也可以考虑用大 batch_size 了。
