苏剑林作家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

引荐、解读、议论和报道人工智能前沿论文效果的学术平台。

表面梯度下降NLPKeras
1
暂无评论
暂无评论~