Kerasでメモリ使用量を減らしたかったらmax_queue_sizeを調整しよう

Tweet

Pocket

Kerasで大きめの画像を使ったモデルを訓練していると、メモリが足りなくなるということがよくあります。途中処理の変数のデータ型（np.uint8）を変えるのだけではなく、max_queue_sizeの調整をする必要があることがあります。それを見ていきます。

メモリサイズの目安

ニューラルネットワークに食わせる変数は基本的にfloat32になるので、4バイト変数になります。つまり、「縦解像度×横解像度×チャンネル数（カラーなら3）×バッチサイズ×4バイト」必要になります。y横軸に縦横の解像度、縦軸にバッチサイズを取って計算すると次のような表になります。単位はGBです。

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

Tweet

Pocket