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を下げたほうが効果的かもしれないよ」ということでした。