現(xiàn)在 Keras 中你也可以用小的 batch size 實現(xiàn)大 batch size 的效果了——只要你愿意花 n 倍的時間，可以達(dá)到 n 倍 batch size 的效果，而不需要增加顯存。

作者丨蘇劍林

研究方向丨NLP，神經(jīng)網(wǎng)絡(luò)

個人主頁丨kexue.fm

Github地址：

https://github.com/bojone/accum_optimizer_for_keras

扯淡

在一兩年之前，做 NLP 任務(wù)都不用怎么擔(dān)心 OOM 問題，因為相比 CV 領(lǐng)域的模型，其實大多數(shù) NLP 模型都是很淺的，極少會顯存不足。幸運或者不幸的是，Bert 出世了，然后火了。Bert 及其后來者們（GPT-2、XLNET 等）都是以足夠龐大的 Transformer 模型為基礎(chǔ)，通過足夠多的語料預(yù)訓(xùn)練模型，然后通過 fine tune 的方式來完成特定的 NLP 任務(wù)。?

即使你很不想用 Bert，但現(xiàn)在的實際情況是：你精心設(shè)計的復(fù)雜的模型，效果可能還不如簡單地 fine tune 一下 Bert 好。所以不管怎樣，為了跟上時代，總得需要學(xué)習(xí)一下 Bert 的 fine tune。

問題是“不學(xué)不知道，一學(xué)嚇一跳”，只要任務(wù)稍微復(fù)雜一點，或者句子長度稍微長一點，顯存就不夠用了，batch size 急劇下降——32？16？8？一跌再跌都是有可能的。?

這不難理解，Transformer 基于 Attention，而 Attention 理論上空間和時間復(fù)雜度都是，雖然在算力足夠強的時候，Attention 由于其并行性還是可以表現(xiàn)得足夠快，但是顯存占用量是省不了了，意味著當(dāng)你句子長度變成原來的 2 倍時，顯存占用基本上就需要原來的 4 倍，這個增長比例肯定就容易 OOM 了。?

而更不幸的消息是，大家都在 fine tune 預(yù)訓(xùn)練 Bert 的情況下，你 batch_size=8 可能比別人 batch_size=80 低好幾個千分點甚至是幾個百分點，顯然這對于要刷榜的讀者是很難受的。難道除了加顯卡就沒有別的辦法了嗎？

正事

有！通過梯度緩存和累積的方式，用時間來換取空間，最終訓(xùn)練效果等效于更大的 batch size。因此，只要你跑得起 batch_size=1，只要你愿意花 n 倍的時間，就可以跑出 n 倍的 batch size 了。?

梯度累積的思路，在之前的文章“讓Keras更酷一些！”：小眾的自定義優(yōu)化器已經(jīng)介紹了，當(dāng)時稱之為“軟 batch（soft batch）”，本文還是沿著主流的叫法稱之為“梯度累積（accumulate gradients）”好了。

所謂梯度累積，其實很簡單，我們梯度下降所用的梯度，實際上是多個樣本算出來的梯度的平均值，以 batch_size=128 為例，你可以一次性算出 128 個樣本的梯度然后平均，我也可以每次算 16 個樣本的平均梯度，然后緩存累加起來，算夠了 8 次之后，然后把總梯度除以 8，然后才執(zhí)行參數(shù)更新。當(dāng)然，必須累積到了 8 次之后，用 8 次的平均梯度才去更新參數(shù)，不能每算 16 個就去更新一次，不然就是 batch_size=16 了。?

剛才說了，在之前的文章的那個寫法是有誤的，因為用到了：

K.switch(cond,?K.update(p,?new_p),?p)

來控制更新，但事實上這個寫法不能控制更新，因為 K.switch 只保證結(jié)果的選擇性，不保證執(zhí)行的選擇性，事實上它等價于：

cond?*?K.update(p,?new_p)?+?(1?-?cond)?*?p

也就是說不管 cond 如何，兩個分支都是被執(zhí)行了。事實上 Keras 或 Tensorflow“幾乎”不存在只執(zhí)行一個分支的條件寫法（說“幾乎”是因為在一些比較苛刻的條件下可以做到），所以此路不通。

不能這樣寫的話，那只能在“更新量”上面下功夫，如前面所言，每次算 16 個樣本的梯度，每次都更新參數(shù)，只不過 8 次中有 7 次的更新量是 0，而只有 1 次是真正的梯度下降更新。

很幸運的是，這種寫法還可以無縫地接入到現(xiàn)有的 Keras 優(yōu)化器中，使得我們不需要重寫優(yōu)化器！詳細(xì)寫法請看：

https://github.com/bojone/accum_optimizer_for_keras

具體的寫法無外乎就是一些移花接木的編程技巧，真正有技術(shù)含量的部分不多。關(guān)于寫法本身不再細(xì)講，如果有疑問歡迎討論區(qū)討論。?

注：這個優(yōu)化器的修改，使得小 batch size 能起到大 batch size 的效果，前提是模型不包含 Batch Normalization，因為 Batch Normalization 在梯度下降的時候必須用整個 batch 的均值方差。所以如果你的網(wǎng)絡(luò)用到了 Batch Normalization，想要準(zhǔn)確達(dá)到大 batch size 的效果，目前唯一的方法就是加顯存/加顯卡。

實驗

至于用法則很簡單：

opt?=?AccumOptimizer(Adam(),?10)?#?10是累積步數(shù)model.compile(loss='mse',?optimizer=opt)model.fit(x_train,?y_train,?epochs=10,?batch_size=10)#?10是累積步數(shù)
model.compile(loss='mse',?optimizer=opt)
model.fit(x_train,?y_train,?epochs=10,?batch_size=10)

這樣一來就等價于 batch_size=100 的 Adam 優(yōu)化器了，代價就是你跑了 10 個 epoch，實際上只相當(dāng)于 batch_size=100 跑了 1 個 epoch，好處是你只需要用到 batch_size=10 的顯存量。?

可能讀者想問的一個問題是：你怎么證明你的寫法生效了？也就是說你怎么證明你的結(jié)果確實是 batch_size=100 而不是 batch_size=10？

為此，我做了個比較極端的實驗，代碼在這里：

https://github.com/bojone/accum_optimizer_for_keras/blob/master/mnist_mlp_example.py?

代碼很簡單，就是用多層 MLP 做 MNIST 分類，用 Adam 優(yōu)化器， fit 的時候 batch_size=1。優(yōu)化器有兩個選擇，第一個是直接 Adam() ，第二個是 AccumOptimizer(Adam(), 100) ：

如果是直接 Adam() ，那 loss 一直在 0.4 上下徘徊，后面 loss 越來越大了（訓(xùn)練集都這樣），val 的準(zhǔn)確率也沒超過 97%；?

如果是 AccumOptimizer(Adam(), 100) ，那么訓(xùn)練集的 loss 越來越低，最終降到 0.02 左右，val 的最高準(zhǔn)確率有 98%+；?

最后我比較了直接 Adam() 但是 batch_size=100 的結(jié)果，發(fā)現(xiàn)跟 AccumOptimizer(Adam(), 100) 但是 batch_size=1 時表現(xiàn)差不多。?

這個結(jié)果足以表明寫法生效了，達(dá)到了預(yù)期的目的。如果這還不夠說服力，我再提供一個訓(xùn)練結(jié)果作為參考：

在某個 Bert 的 fine tune 實驗中，直接用 Adam() 加 batch_size=12，我跑到了 70.33% 的準(zhǔn)確率；我用 AccumOptimizer(Adam(), 10) 加 batch_size=12（預(yù)期等效 batch size 是 120），我跑到了 71% 的準(zhǔn)確率，提高了 0.7%，如果你在刷榜，那么這 0.7% 可能是決定性的。

結(jié)論

終于把梯度累積（軟 batch）正式地實現(xiàn)了，以后用 Bert 的時候，也可以考慮用大 batch_size 了。

點擊以下標(biāo)題查看作者其他文章：?

• 變分自編碼器VAE：原來是這么一回事 | 附源碼

• 再談變分自編碼器VAE：從貝葉斯觀點出發(fā)

• 變分自編碼器VAE：這樣做為什么能成？

• 簡單修改，讓GAN的判別器秒變編碼器
  

• 深度學(xué)習(xí)中的互信息：無監(jiān)督提取特征

• 全新視角：用變分推斷統(tǒng)一理解生成模型

• 細(xì)水長flow之NICE：流模型的基本概念與實現(xiàn)

• 細(xì)水長flow之f-VAEs：Glow與VAEs的聯(lián)姻

• 深度學(xué)習(xí)中的Lipschitz約束：泛化與生成模型
  

#好 書 推 薦#

?深度學(xué)習(xí)理論與實戰(zhàn)：基礎(chǔ)篇?

李理 / 編著

本書不僅包含人工智能、機器學(xué)習(xí)及深度學(xué)習(xí)的基礎(chǔ)知識，如卷積神經(jīng)網(wǎng)絡(luò)、循環(huán)神經(jīng)網(wǎng)絡(luò)、生成對抗網(wǎng)絡(luò)等，而且也囊括了學(xué)會使用 TensorFlow、PyTorch 和 Keras 這三個主流的深度學(xué)習(xí)框架的*小知識量；不僅有針對相關(guān)理論的深入解釋，而且也有實用的技巧，包括常見的優(yōu)化技巧、使用多 GPU 訓(xùn)練、調(diào)試程序及將模型上線到生產(chǎn)系統(tǒng)中。

本書希望同時兼顧理論和實戰(zhàn)，使讀者既能深入理解理論知識，又能把理論知識用于實戰(zhàn)，因此本書每介紹完一個模型都會介紹其實現(xiàn)，讀者閱讀完一個模型的介紹之后就可以運行、閱讀和修改相關(guān)代碼，從而可以更加深刻地理解理論知識。

?長按識別二維碼查看詳情?

?

現(xiàn)在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關(guān)注」訂閱我們的專欄吧

關(guān)于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學(xué)術(shù)平臺。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

▽ 點擊 |?閱讀原文?| 查看作者博客

總結(jié)

以上是生活随笔為你收集整理的Keras梯度累积优化器：用时间换取效果的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。