Kerasでメモリ使用量を減らしたかったらmax_queue_sizeを調整しよう
Kerasで大きめの画像を使ったモデルを訓練していると、メモリが足りなくなるということがよくあります。途中処理の変数のデータ型（np.uint8）を変えるのだけではなく、max_queue_sizeの調整をする必要があることがあります。それを見ていきます。
メモリサイズの目安
ニューラルネットワークに食わせる変数は基本的にfloat32になるので、4バイト変数になります。つまり、「縦解像度×横解像度×チャンネル数（カラーなら3）×バッチサイズ×4バイト」必要になります。y横軸に縦横の解像度、縦軸にバッチサイズを取って計算すると次のような表になります。単位はGBです。
| バッチ/解像度 | 32 | 64 | 128 | 256 | 512 | 1024 |
|---|---|---|---|---|---|---|
| 32 | 0.000366211 | 0.001464844 | 0.005859375 | 0.0234375 | 0.09375 | 0.375 |
| 64 | 0.000732422 | 0.002929688 | 0.01171875 | 0.046875 | 0.1875 | 0.75 |
| 128 | 0.001464844 | 0.005859375 | 0.0234375 | 0.09375 | 0.375 | 1.5 |
| 256 | 0.002929688 | 0.01171875 | 0.046875 | 0.1875 | 0.75 | 3 |
| 512 | 0.005859375 | 0.0234375 | 0.09375 | 0.375 | 1.5 | 6 |
| 1024 | 0.01171875 | 0.046875 | 0.1875 | 0.75 | 3 | 12 |
| 2048 | 0.0234375 | 0.09375 | 0.375 | 1.5 | 6 | 24 |
| 4096 | 0.046875 | 0.1875 | 0.75 | 3 | 12 | 48 |
| 8192 | 0.09375 | 0.375 | 1.5 | 6 | 24 | 96 |
| 16384 | 0.1875 | 0.75 | 3 | 12 | 48 | 192 |
| 32768 | 0.375 | 1.5 | 6 | 24 | 96 | 384 |
| 65536 | 0.75 | 3 | 12 | 48 | 192 | 768 |
32×32ぐらいではほぼ問題になることはありませんが、256×256ぐらいから結構問題になってくるのではないでしょうか。
ただ、これは最小限の値で、実際にはこれにラベルデータや、主にバッチの演算や代入をする場合は瞬間的にこれの2倍ぐらいは見ておいたほうがいいのではないかと思います。ここに出るくる値が2GBでも、実際は10GBぐらい消費していたなんてことがあります。
1つの大きな原因はmax_queue_size
すべてがすべて解決できるわけではないですが、実際のメモリ使用量を減らす方策として、model.fit_generator, predict_generatorにおける「max_queue_size」を減らすというのがあります。これを調整している人はあまり見たことない気がします。
公式ドキュメントを見ると、model.fit_generatorの中に
max_queue_size: 整数．ジェネレータのキューのための最大サイズ． 指定しなければmax_queue_sizeはデフォルトで10になります．
特に指定しないと、デフォルトで10もキュー（キャッシュ）しているんですね。この単位はバッチなので、10バッチ分キャッシュしていると考えて良いと思います。GPUの性能があまり良くなくて、データの読み込みよりもネットワークの演算のほうがボトルネックになりやすいケースではこれは有効です。しかし、最近のGPUや特にTPUのように、ネットワークの演算は高性能だけど、データの読み込みが追いつかなくてジェネレーターがボトルネックになるケースでは、このキャッシュがほとんど意味をなさなくなります。
まずmax_queue_size分キャッシュしているのを確認します。
from keras.layers import Dense, Input
from keras.models import Model
import numpy as np
from keras.utils import to_categorical
input = Input((60,))
x = Dense(10, activation="softmax")(input)
model = Model(input, x)
model.compile("adam", "categorical_crossentropy", ["acc"])
def base_generator():
# batch_size=128
while True:
X = np.random.randn(128,60)
y = to_categorical(np.random.randint(0, 10,128), num_classes=10)
yield X, y
y_cache = []
def cache_generator():
global y_cache
y_cache = np.zeros(10).reshape(1,-1)
print("reset cache")
for X, y in base_generator():
y_cache = np.concatenate([y_cache, y], axis=0)
yield X, y
y_pred = model.predict_generator(cache_generator(), steps=10, max_queue_size=10)
print("Generatorが実際に転送した値")
print(y_pred.shape)
print("内部キャッシュ")
cache = np.asarray(y_cache[1:])
print(cache.shape)
このように、base_generator()をラップするようなcache_generator()を作り、ここでYの値をキャッシュしながらpredict_generatorさせています。ここでのキャッシュしたYの値は、実際にニューラルネットワーク側に転送されなくても、キュー用に内部的に読み込んだデータ数になるので、predictで返したサンプル数よりも多くなります。実際、デフォルトの（max_queue_size=10）で実行すると、
Generatorが実際に転送した値
(1280, 10)
内部キャッシュ
(2688, 10)
このように実際に転送したサンプルの2倍以上キューしていることになります。ちなみにこのmax_queue_sizeを2にすると、
Generatorが実際に転送した値
(1280, 10)
内部キャッシュ
(1536, 10)
キャッシュのマージンが少なくなります。
データの読み込みがボトルネックになっているケースでは、10バッチもキャッシュする必要はありません。仮にメモリが溢れて、その対策としてバッチサイズを下げ、ネットワークの計算が遅くなる/安定性が悪くなるのだったら、max_queue_sizeを下げたほうが効果があると思います。あくまで読み込みがボトルネックになるケースではですが。
逆に読み込みがボトルネックにならないようなケースでは、このデフォルトの設定はかなり有効になります。計算デバイスの性能良くなって、どこに重きをおくかが変わってきたのでしょうね。
まとめ
「メモリサイズをケチりたかったら、バッチサイズを下げるだけじゃなくて、max_queue_sizeを下げたほうが効果的かもしれないよ」ということでした。