现在 Keras 中你也可以用小的 batch size 实现大 batch size 的效果了——只要你愿意花 n 倍的时间，可以达到 n 倍 batch size 的效果，而不需要增加显存。

作者丨苏剑林

研究方向丨NLP，神经网络

个人主页丨kexue.fm

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)# 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 了。

点击以下标题查看作者其他文章：

• 变分自编码器VAE：原来是这么一回事 | 附源码

• 再谈变分自编码器VAE：从贝叶斯观点出发

• 变分自编码器VAE：这样做为什么能成？

• 简单修改，让GAN的判别器秒变编码器

• 深度学习中的互信息：无监督提取特征

• 全新视角：用变分推断统一理解生成模型

• 细水长flow之NICE：流模型的基本概念与实现

• 细水长flow之f-VAEs：Glow与VAEs的联姻

• 深度学习中的Lipschitz约束：泛化与生成模型

#好 书 推 荐#

 深度学习理论与实战：基础篇 

李理 / 编著

本书不仅包含人工智能、机器学习及深度学习的基础知识，如卷积神经网络、循环神经网络、生成对抗网络等，而且也囊括了学会使用 TensorFlow、PyTorch 和 Keras 这三个主流的深度学习框架的*小知识量；不仅有针对相关理论的深入解释，而且也有实用的技巧，包括常见的优化技巧、使用多 GPU 训练、调试程序及将模型上线到生产系统中。

本书希望同时兼顾理论和实战，使读者既能深入理解理论知识，又能把理论知识用于实战，因此本书每介绍完一个模型都会介绍其实现，读者阅读完一个模型的介绍之后就可以运行、阅读和修改相关代码，从而可以更加深刻地理解理论知识。

 长按识别二维码查看详情 

?

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

关于PaperWeekly

PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域，欢迎在公众号后台点击「交流群」，小助手将把你带入 PaperWeekly 的交流群里。

▽ 点击 | 阅读原文 | 查看作者博客