Optuna + Keras が GPU のメモリを食いつぶす

Optuna で Keras (TensorFlow GPU) のハイパーパラメータを最適化しようと思ったのですが、しばらく trial を繰り返すと GPU の OOM エラーが発生し失敗してしまいました。解決までかなり手こずったので、忘れないうちにメモしておきます。

なお、私がこの問題に遭遇したのがたまたま Optana で最適化しようとしたタイミングだったというだけで、問題自体は Optana に限った話ではありません。クロスバリデーションやグリッドサーチなど、1回の実行で複数回 Keras の学習処理が GPU で走るような処理を書いた場合に常に起こりえます。

メモリを食いつぶすコード例 (Before)

Optuna デビュー戦ということで試しに学習率を最適化してみようと思い、以下のようなコードを書きました (※説明に必要な部分だけ抜粋したため動かないかも……)。各 trial でデータは共通なので、あらかじめ読み込んでおき、objective 内ではそれを参照するだけにしました。これが原因で OOM 地獄に落ちることなど、この時の私は知る由もなかったのです。

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import SGD
import optuna

# データを読み込む
data = pd.read_csv("data.csv")
x_train = data.drop(columns = ["y"]).values
y_train = data["y"]

# optunaで最適化する目的関数を定義する。ここでは学習率を最適化する
def objective(trial):
    lr = trial.suggest_loguniform("lr", 1e-5, 1e-1)
    
    model = Sequential()
    model.add(Dense(units = 128, activation = "relu"))
    model.add(Dense(units = 128, activation = "relu"))
    model.add(Dense(units = 128, activation = "relu"))
    model.add(Dense(1))
    
    optimizer = SGD(lr = lr)
    model.compile(loss = "mean_squared_error", optimizer = optimizer)
    history = model.fit(x = x_train, y = y_train, epochs = 100, validation_split = 0.2, verbose = 0)
    
    return -np.amax(history.history['val_acc'])

# 最適化する
study = optuna.create_study()
study.optimize(objective, n_trials = 200)

メモリを食いつぶさないコード例 (After)

私は before のコードを実行し、意気揚々と会社に向かいました。しかし帰宅すると、そこには大量の OOM エラーで真っ赤に染まった無残な jupyter notebook の姿が……！

悲しみに暮れながらも、頑張って "keras gpu memory release" でググります。すると「clear_session() すれば解放できるよ」という情報が見つかったため試してみたものの、解放されませんでした。さらに調べると「明示的に del model するといいよ」「俺は del history も必要だったね」という情報も出てきましたが、やはり解放されませんでした。

仕方がないのでググるのをやめ、自分のコードとOptuna 公式の Keras のサンプルコードを見比べ違いを探しました。サンプルコードではデータを objective の中で読み込んでいるんですね。ここが自分のコードとの最大の違いでした。これは想像ですが、サンプルコードでは trial のたびに objective のスコープ内でデータを読み込むため、trial が終了するたびに (objective のスコープを抜けるたびに) データが GPU のメモリ上から解放され、問題なく動くのでしょう。一方で、私の書いたコードでは objective の外で読み込んだデータを参照しているため、trial が終わってもデータを解放できず、trial を繰り返すたびにゴミが溜まっていくのだと考えました。

冷静に考えれば、モデルや履歴のサイズなどデータと比べれば高が知れているため del model del history する意味はあまりありません。del すべきはデータです。そこで、コードを以下のように書き換えました。

1. データを各 trial 用にコピーする
2. コピーしたデータを明示的に破棄する

このように書き換えたところ、無事メモリを食いつぶすことなく最適化を実行できました。めでたしめでたし。

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import SGD
from keras.backend import clear_session
import optuna

# データを読み込む
data = pd.read_csv("data.csv")
x_train = data.drop(columns = ["y"]).values
y_train = data["y"]

# optunaで最適化する目的関数を定義する。ここでは学習率を最適化する
def objective(trial):
    lr = trial.suggest_loguniform("lr", 1e-5, 1e-1)
    
    # 変更点1: データを各trial用にコピーする
    x_train_copy = np.copy(x_train)
    y_train_copy = np.copy(y_train)
    
    model = Sequential()
    model.add(Dense(units = 128, activation = "relu"))
    model.add(Dense(units = 128, activation = "relu"))
    model.add(Dense(units = 128, activation = "relu"))
    model.add(Dense(1))
    
    optimizer = SGD(lr = lr)
    model.compile(loss = "mean_squared_error", optimizer = optimizer)
    history = model.fit(x = x_train_copy, y = y_train_copy, epochs = 100, validation_split = 0.2, verbose = 0)
    
    # 変更点2: コピーしたデータを明示的に破棄する
    clear_session()
    del model, optimizer, history, x_train_copy, y_train_copy
    gc.collect()
    
    return -np.amax(history.history['val_acc'])

# 最適化する
study = optuna.create_study()
study.optimize(objective, n_trials = 200)