KerasがGPUメモリを食いつぶしてしまう

Question

Tensorflow-gpu(ver2.1)とKeras-gpu(ver2.3.1)を用いてJyupter notebook上で画像分類学習をしているのですが、学習結果を用いて分類させている際に数回分類させた後、

ResourceExhaustedError: OOM when allocating tensor with shape[712320,128] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:Add] name: dense_1_15/random_uniform/

とエラーが出てプログラムが止まってしまいます。

タスクマネージャーを監視していると、学習結果を用いる際にGPUメモリを0.3GBから7.0GBまで確保しており、KerasがGPUメモリを食いつぶした結果プログラムが停止するのかなと考えています。

対策法を調べてみると、

physical_devices = tf.config.list_physical_devices('GPU')
if len(physical_devices) > 0:
    for device in physical_devices:
        tf.config.experimental.set_memory_growth(device, True)
        print('{} memory growth: {}'.format(device, tf.config.experimental.get_memory_growth(device)))
else:
    print("Not enough GPU hardware devices available")

こちらでGPUメモリの利用を制限できるそうなのですが、これをコードの先頭付近にいれて実行しても変わらず7.0GBまで確保してしまい、結局エラーがでます。

どなたか解決策をご存知でしょうか？
よろしくお願いいたします。

Answer 1

各notebookの先頭に下記を追加したら、必要な分だけしかメモリーが使われなくなるはずです

import os
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

ただし、一旦確保されたメモリーはpythonが完全に落ちるまでは開放されません
AIの実行が終わっても、AIの実行に使ったpythonが残っているとメモリーが確保されたままになるので、メモリー不足になりやすくなります
必要無くなったpythonは、こまめに完全終了させることをお勧めします

Webブラウザのタブを閉じるだけではpythonは終了せず、メモリーは開放されないようです
cuDNN初期化エラーが出てTensorFlowGPUの学習ができなくなったけど解決した話
JupyterNotebookでGPUのメモリを解放する方法-Windows編