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

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

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

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

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

```py
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()` すれば解放できるよ](https://stackoverflow.com/questions/43930871/how-to-release-the-occupied-gpu-memory-when-calling-keras-model-by-apache-mod-ws)」という情報が見つかったため試してみたものの、解放されませんでした。さらに調べると「[明示的に `del model` するといいよ](https://forums.fast.ai/t/how-could-i-release-gpu-memory-of-keras/2023/4)」「[俺は `del history` も必要だったね](https://forums.fast.ai/t/how-could-i-release-gpu-memory-of-keras/2023/7)」という情報も出てきましたが、やはり解放されませんでした。

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

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

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

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

```py
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)
```