こしあん
2018-12-28

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

Related Posts

note開設のお知らせ 本日noteを開設いたしました。 https://note.mu/koshian2 これは自分の記事をより多くの方々に読んでいただき、新たな読者の開拓を図るためであります。 当面は既存の記事の再送を中心に考えていますが、いくつかnote向けに読みやすい新規の記事も考えています。好評なら新規の...
Kerasのジェネレーターでサンプルが列挙される順番について... Kerasの(カスタム)ジェネレーターでサンプルがどの順番で呼び出されるか、1ループ終わったあとにどういう処理がなされるのか調べてみました。ジェネレーターを自分で定義するとモデルの表現の幅は広がるものの、バグが起きやすくなるので「本当に順番が保証されるのか」や「ハマりどころ」を確認します。 0~...
TPUで学習率減衰させる方法 TPUで学習率減衰したいが、TensorFlowのオプティマイザーを使うべきか、tf.kerasのオプティマイザーを使うべきか、あるいはKerasのオプティマイザーを使うべきか非常にややこしいことがあります。TPUで学習率を減衰させる方法を再現しました。 結論から TPU環境でtf.keras...
Kerasで損失関数に複数の変数を渡す方法... Kerasで少し複雑なモデルを訓練させるときに、損失関数にy_true, y_pred以外の値を渡したいときがあります。クラスのインスタンス変数などでキャッシュさせることなく、ダイレクトに損失関数に複数の値を渡す方法を紹介します。 元ネタ:Passing additional arguments...
Python(Numpy)で画像を水平反転する方法:Data Augmentation向け... OpenCVを使わずに単純に画像を左右反転(水平反転)する方法を考えます。ディープラーニングでデータのジェネレーターを自分で実装した場合、Data Augmentationを組み込む際にも必要になります。それを見ていきましょう。 左右反転自体は実は簡単 例えばNumpyの行列を左右反転させてみ...

Add a Comment

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です