こしあん
2018-12-28

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

Pocket
LINEで送る


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

PyTorch/TorchVisionで複数の入力をモデルに渡したいケース... PyTorch/TorchVisionで入力が複数あり、それぞれの入力に対して同じ前処理(transforms)をかけるケースを考えます。デフォルトのtransformsは複数対応していないのでうまくいきません。しかし、ラッパークラスを作り、それで前処理をラップするといい感じにできたのでその方法を...
pix2pix HDのCoarse to fineジェネレーターを考える... pix2pix HDの論文を読んでいたら「Coarse to fineジェネレーター」という、低解像度→高解像度と解像度を分けて訓練するネットワークの工夫をしていました。pix2pixはGANですが、このジェネレーターや訓練の工夫は、Non-GANでも理屈上は使えるはずなので、この有効性をImag...
Chainerで画像の前処理やDataAugmentationをしたいときはDatasetMixin... Chainerにはデフォルトでランダムクロップや標準化といった、画像の前処理やDataAugmentation用の関数が用意されていません。別途のChainer CVというライブラリを使う方法もありますが、chainer.dataset.DatasetMixinを継承させて独自のデータ・セットを定...
画像のダウンサンプリングとPSNRの測定... U-Netでどこまでの深さが効いているのかを考えるために、画像をダウンサンプリングし、アップサンプリングするという処理を行いPSNRを見ていきます。その結果、PSNRが15~20程度だと、U-Netの深い層が効かないかもしれないという状況を確認することができました。 きかっけ・考え方 U-Ne...
WarmupとData Augmentationのバッチサイズ別の精度低下について... 大きいバッチサイズで訓練する際は、バッチサイズの増加にともなう精度低下が深刻になります。この精度低下を抑制することはできるのですが、例えばData Augmentationのようなデータ増強・正則化による精度向上とは何が違うのでしょうか。それを調べてみました。 きっかけ この記事を書いたときに...
Pocket
Delicious にシェア

Add a Comment

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