【Python】Tensorflow/Kerasのループ学習でメモリ不足の解決方法 : すきま研究所日誌

Tensorflow / Kerasを用いた深層学習のループ実行でメモリ使用量が増え続け、メモリオーバーフローが発生する問題が発生しました。

調査して対応したことを備忘録としてまとめておきます。

Tensorflowループでメモリリーク
メモリ使用量増加と原因考察
メモリ増加の解決策

Tensorflowループでメモリリーク

以下のTensorFlow / Kerasを用いた訓練を行う関数をループした時、メモリ使用量が増え続けオーバーフローしました。

Python（ループコード）：

for i in range(loop_num):
    # tensorflow                
    model, scaler = deep_learning_model_train.deep_learning_model_train(df_x, df_t)

Python（tensorflow本体）：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, BatchNormalization, Dropout, Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping
from sklearn.preprocessing import StandardScaler
from tensorflow.keras import backend as K

def deep_learning_model_train(x_train, t_train):
    # 標準化
    scaler_y = StandardScaler()
    t_train = scaler_y.fit_transform(t_train.to_numpy().reshape(-1, 1))
    # モデル構築
    model = Sequential([
        Input(shape=(x_train.shape[1],)),
        Dense(128, activation='relu'),
        BatchNormalization(),
        Dropout(0.2),
        Dense(64, activation='relu'),
        BatchNormalization(),
        Dropout(0.2),
        Dense(32, activation='relu'),
        Dense(1)
    ])
    # コンパイル
    model.compile(optimizer=Adam(learning_rate=0.001), loss='mean_squared_error', metrics=['mse'])
    # 早期停止
    early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)   
    # 学習
    model.fit(
        x_train, t_train,
        epochs=200, batch_size=32,
        validation_split=0.1, verbose=0,
        callbacks=[early_stopping]
    )
    return model, scaler_y

メモリ使用量増加と原因考察

メモリ使用量確認

次の処理をコード中に追加して、メモリが増加する様子をLogへ出力してみます。

Python（メモリ使用量確認）：

for i in range(50):
    # tensorflow                
    model, scaler = deep_learning_model_train.deep_learning_model_train(df_x, df_t)
    print(f"After processing: {get_memory_usage.get_memory_usage() / 1024 / 1024:.2f} MB")

▼関連記事

出力 Log：

After processing: 577.95 MB
After processing: 609.11 MB
After processing: 636.73 MB
After processing: 636.13 MB
After processing: 662.74 MB
After processing: 692.16 MB
After processing: 725.83 MB
After processing: 703.88 MB
After processing: 725.17 MB
After processing: 755.32 MB

10回ほどループを回しただけでもメモリ使用量がみるみる増えていることが分かります。

1回の訓練あたり20MBと少し増加といったところでしょうか

メモリ増加発生要因考察

ひとまずそれらしいキーワードで検索してみると以下が見つかりました。

参考：
https://kurupical.hatenablog.com/entry/2017/12/02/233735

計算グラフのキャッシュが蓄積していくことで最終的にメモリ不足となり、実行でKilled（OSによる強制停止）を引き起こしているようです。

計算グラフは、誤解を恐れずに言うと他の環境や言語に移植したり、並列処理などにより処理の高速化を狙ったりするために生成されているもののため、訓練自体の目的には特に必要ありません。

メモリ増加の解決策

以下のようにコードに処理を追加して、計算グラフの発生抑制かつ計算グラフのリセットをしてみます。

Python (tensorflow)：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, BatchNormalization, Dropout, Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping
from sklearn.preprocessing import StandardScaler
from tensorflow.keras import backend as K

def deep_learning_model_train(x_train, t_train):
    # 標準化
    scaler_y = StandardScaler()
    t_train = scaler_y.fit_transform(t_train.to_numpy().reshape(-1, 1))
    # Eager Execution有効化（デバッグモード）
    tf.data.experimental.enable_debug_mode()
    # 即時実行（Eager Execution）有効化 (計算グラフを作成しない)
    tf.config.run_functions_eagerly(True)
    # モデル構築
    model = Sequential([
        Input(shape=(x_train.shape[1],)),
        Dense(128, activation='relu'),
        BatchNormalization(),
        Dropout(0.2),
        Dense(64, activation='relu'),
        BatchNormalization(),
        Dropout(0.2),
        Dense(32, activation='relu'),
        Dense(1)
    ])
    # コンパイル
    model.compile(optimizer=Adam(learning_rate=0.001), loss='mean_squared_error', metrics=['mse'])
    # 早期停止
    early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)   
    # 学習
    model.fit(
        x_train, t_train,
        epochs=200, batch_size=32,
        validation_split=0.1, verbose=0,
        callbacks=[early_stopping]
    )
    # メモリ解放
    K.clear_session()
    tf.compat.v1.reset_default_graph()
    return model, scaler_y

追加したのは、次の２か所です。

13~16行目：

tf.data.experimental.enable_debug_mode()を使用してデバッグモード
tf.config.run_functions_eagerly(True)を使用し、計算グラフを作らない（Eager Execution）で実行

41~42行目：

セッションのクリア
グラフのリセット（TensorFlow v1 用？）

出力Log メモリ使用量確認：

After processing: 540.20 MB
After processing: 541.18 MB
After processing: 542.00 MB
After processing: 542.33 MB
After processing: 542.52 MB
After processing: 542.55 MB
After processing: 542.71 MB
After processing: 543.16 MB
After processing: 543.10 MB
After processing: 543.36 MB
After processing: 543.72 MB

異常なメモリ使用量の増加はなくなりました

動作環境確認

Pythonのバージョンと、TensorFlowのバージョン違いの環境では、訓練後の「K.clear_session()」と「tf.compat.v1.reset_default_graph()」のメモリ解放のみでメモリ使用量の増加を抑制することが可能です。

以下の環境では、今回のように「tf.data.experimental.enable_debug_mode()」と「tf.config.run_functions_eagerly(True)」が必要でした。

環境確認↓

コマンド：