苏剑林作者NLP,神经网络研究方向

Keras梯度累积优化器:用时间换取效果

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

PaperWeekly
PaperWeekly

推荐、解读、讨论和报道人工智能前沿论文成果的学术平台。

理论KerasNLP梯度下降
1
暂无评论
暂无评论~