前言

在前不久的文章《BERT重計算：用22.5%的訓練時間節省5倍的視訊記憶體開銷（附程式碼）》中介紹了一個叫做“重計算”的技巧（附pytorch和paddlepaddle實現）。簡單來說重計算就是用來省視訊記憶體的方法，讓平均訓練速度慢一點，但batch_size可以增大好幾倍，該技巧首先釋出於論文《Training Deep Nets with Sublinear Memory Cost》。

最近筆者發現，重計算的技巧在tensorflow也有實現。事實上從tensorflow1.8開始，tensorflow就已經自帶了該功能了，當時被列入了tf.contrib這個子庫中，而從tensorflow1.15開始，它就被內建為tensorflow的主函式之一，那就是tf.recompute_grad。找到 tf.recompute_grad 之後，筆者就琢磨了一下它的用法，經過一番折騰，最終居然真的成功地用起來了，居然成功地讓 batch_size 從48增加到了144！然而，在繼續整理測試的過程中，發現這玩意居然在tensorflow 2.x是失效的...於是再折騰了兩天，查詢了各種資料並反覆除錯，最終算是成功地補充了這一缺陷。

最後是筆者自己的開源實現：

Github地址：

https://github.com/bojone/keras_recompute

該實現已經內建在bert4keras中，使用bert4keras的讀者可以升級到最新版本（0.7.5+）來測試該功能。

使用

筆者的實現也命名為recompute_grad，它是一個裝飾器，用於自定義Keras層的 call函式，比如

from recompute import recompute_grad
class MyLayer(Layer):@recompute_graddef call(self, inputs):return inputs * 2

對於已經存在的層，可以透過繼承的方式來裝飾：

from recompute import recompute_grad
class MyDense(Dense):@recompute_graddef call(self, inputs):return super(MyDense, self).call(inputs)

自定義好層之後，在程式碼中嵌入自定義層，然後在執行程式碼之前，加入環境變數RECOMPUTE=1來啟用重計算。

注意：不是在總模型裡插入了@recomputr_grad，就能達到省記憶體的目的，而是要在每個層都插入@recomputr_grad才能更好地省視訊記憶體。簡單來說，就是插入的@recomputr_grad越多，就省視訊記憶體。具體原因請仔細理解重計算的原理。

效果

bert4keras0.7.5已經內建了重計算，直接傳入環境變數RECOMPUTE=1就會啟用重計算，讀者可以自行嘗試，大概的效果是：

1、在BERT Base版本下，batch_size可以增大為原來的3倍左右；

2、在BERT Large版本下，batch_size可以增大為原來的4倍左右；

3、平均每個樣本的訓練時間大約增加25%；

4、理論上，層數越多，batch_size可以增大的倍數越大。

環境

在下面的環境下測試透過：

tensorflow 1.14 + keras 2.3.1

tensorflow 1.15 + keras 2.3.1

tensorflow 2.0 + keras 2.3.1

tensorflow 2.1 + keras 2.3.1

tensorflow 2.0 + 自帶tf.keras

tensorflow 2.1 + 自帶tf.keras

確認不支援的環境：

tensorflow 1.x + 自帶tf.keras

歡迎報告更多的測試結果。

順便說一下，強烈建議用keras2.3.1配合tensorflow1.x/2.x來跑，強烈不建議使用tensorflow 2.x自帶的tf.keras來跑。