こしあん
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

Chainerで画像の前処理やDataAugmentationをしたいときはDatasetMixin... Chainerにはデフォルトでランダムクロップや標準化といった、画像の前処理やDataAugmentation用の関数が用意されていません。別途のChainer CVというライブラリを使う方法もありますが、chainer.dataset.DatasetMixinを継承させて独自のデータ・セットを定...
Google ColabのTPU環境でmodel.fitのhistoryが消える現象... Google ColabのTPU環境でmodel.fitしたときに、通常の環境で得られるhistory(誤差や精度のログ)が消えていることがあります。その対応法を示します。 原因はTPU用のモデルに変換したから まず結論からいうとこの現象はCPU/GPU環境では再発しません。TPU環境特有の現...
KerasのCallbackを使って継承したImageDataGeneratorに値が渡せるか確かめ... Kerasで前処理の内容をエポックごとに変えたいというケースがたまにあります。これを実装するとなると、CallbackからGeneratorに値を渡すというコードになりますが、これが本当にできるかどうか確かめてみました。 想定する状況 例えば、前処理で正則化に関係するData Augmenta...
Kerasのジェネレーターでサンプルが列挙される順番について... Kerasの(カスタム)ジェネレーターでサンプルがどの順番で呼び出されるか、1ループ終わったあとにどういう処理がなされるのか調べてみました。ジェネレーターを自分で定義するとモデルの表現の幅は広がるものの、バグが起きやすくなるので「本当に順番が保証されるのか」や「ハマりどころ」を確認します。 0~...
Python(Numpy)で画像を水平反転する方法:Data Augmentation向け... OpenCVを使わずに単純に画像を左右反転(水平反転)する方法を考えます。ディープラーニングでデータのジェネレーターを自分で実装した場合、Data Augmentationを組み込む際にも必要になります。それを見ていきましょう。 左右反転自体は実は簡単 例えばNumpyの行列を左右反転させてみ...

Add a Comment

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