Keras梯度累积优化器：用时间换取效果
現在 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 如何，兩個分支都是被執行了。事實上 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 的效果，目前唯一的方法就是加顯存/加顯卡。
實驗
至于用法則很簡單：
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 的交流群里。
▽ 點擊 |?閱讀原文?| 查看作者博客
總結
以上是生活随笔為你收集整理的Keras梯度累积优化器：用时间换取效果的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 飞桨PaddlePaddle升级解读 |
- 下一篇： 寻找想改变人工智能的“大人物”！2019