【Ray Tune】Population Based Trainingを使って深層学習のハイパーパラメータをチューニングする with TF2【深層学習】

Population Based Training

PopulationBasedTrainingとは遺伝的アルゴリズムを活かしてニューラルネットワークのハイパーパラメータを最適化するアルゴリズムです。

• みんな大好きDeepMind産
• Grid Searchより探索範囲が少ない
• ほぼ全てのアルゴリズムのハイパラを最適化可能

という心強い性質を備えているので大変人気なアルゴリズムだと思います。
一方でハイパーパラメータの最適化は一般的にいって計算量が膨大なので多くの場合マルチプロセスでやったりクラウド上でマルチノードでやったりします。

しかし分散処理は鬱陶しいのです...。

そこで今回は学習アルゴリズムを入力したらうまいこと勝手に分散処理でハイパラ調整してくれるチューニング専用ライブラリRay Tuneを利用します。

Pythonにおける超クールなマルチプロセッシングライブラリRayについてはこちらで紹介していますのでよかったらこちらも合わせてご覧ください。

Population Based Training とは

Population Based Trainingは早い話が遺伝的アルゴリズムです。
1. ランダムに初期化した１０個のエージェントに並行して学習を進めさせる
2. 学習が終わったらパフォーマンスの高い３個のエージェントを残し残りのエージェントを全て捨てる
3. 捨てられた７個のエージェントに替えてExploreとExploitによって新しいエージェントを生み出す

Explore

ExploreとしてよくあるのはPerturbとResample。
どちらも同時に用いることが可能で、Ray Tuneでは一定確率でランダムに選ばれる仕様です。

Perturbは既存のハイパラに1.2もしくは0.8をランダムに掛け算します。
Resampleはもともとの探索空間からランダムにハイパラを選び出します。

Exploit

Exploitではこれまでのエージェントの中で高いパフォーマンスを残していたモデルをランダムに抽出して全てをコピーしてくるものが多いです。Ray Tuneの実装ではハイパラだけでなく学習した重みもコピーします。

※もちろんここに出てきた１０個とか３個とかそのあたりの具体的な数字は全部適当につけました。

Ray Tuneでの実装

環境

OS: Ubuntu18.04
Python: 3.6.9で動作確認
Deep Learning Framework: Tensorflow2.4 (どんなフレームワークでも可。深層学習でなくてもよき。)

インストール

pip install ray ray[tune]

それでは実装してみましょう。今回はおなじみmnistを学習していきます。ただのMLPって９７％くらいしか精度でないイメージですがハイパラ最適化によって９９％とか出せます。

Import & おまじない

地味に大事なおまじない。

import os

import tensorflow as tf
from tensorflow.keras.datasets import mnist
from ray.tune.integration.keras import TuneReportCheckpointCallback
from ray.tune.schedulers.pbt import PopulationBasedTraining
import ray
from ray import tune

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
def set_growth():
    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)

TF_CPP_MIN_LOG_LEVELを３にしておかないとTensorflowのInfoがコンソールから有益な情報を洗い流してしまいます。
あとはGPU使う人はTensorflowでマルチプロセス学習するときプロセスの開始と同時にこのset_growth()を呼ぶと幸せになれます。

学習のメインループの用意

関数に学習ループをかいて関数ごと渡してやればかってにハイパラを最適化してくれます。
関数はconfigとcheckpoint_dirを引数に取る必要があります。

ハイパラはconfigというdictに入っています。
configからハイパラを読み取って学習する関数を定義してやればよきです。

checkpoint_dirはcheckpoint用のディレクトリのパスで、Exploit用にユーザーが好きな情報を保存することができます。

def train_mnist(config, checkpoint_dir=None):
    # GPU メモリ割り当てがバグらないようになるおまじない (also works for cpu-only machine)
    set growth()

    # batch_size = config["batch_size"]とかけばチューニング対象にできる。
    batch_size = 128
    epochs = 10


    x_train, y_train, x_test, y_test = load_data()
    model = build_model(config, checkpoint_dir)

    model.fit(
        x_train,
        y_train,
        batch_size=batch_size,
        epochs=epochs,
        verbose=0,
        validation_data=(x_test, y_test),
        callbacks=[
            TuneReportCheckpointCallback(
                metrics={"mean_accuracy": "accuracy"},
                filename="model",
                frequency=5,
            ),
        ]
    )

大事なこと

学習時にはパフォーマンスTuneに報告しなければなりません。
これはtune.reportという関数で通常実現されますがTensorflowのFit関数を使って学習する場合はTuneReportCheckpointCallbackを使えばExploit時のためのモデルパラメータの保存もTuneへのパフォーマンスの報告もAutoでやってくれます。TensorflowやTorchなどのメジャーなフレームワークの場合このようなお手軽なAPIが公開されているので探してみてください。

手動で学習する場合は学習ループの中でtune.report(score=score)などとかけばTuneがパフォーマンスを認識してくれます。

あとはトレーニングループの中に使われている関数を実装すれば完成です。
ちなみに別に関数をこんなふうに分類する必要はないです。

データをロードする関数とモデルを用意する関数

def load_data():
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0
    return x_train, y_train, x_test, y_test


def build_model(config: dict, checkpoint_dir: str) -> tf.keras.models.Model:
    num_classes = 10
    if checkpoint_dir:
        checkpoint_filepath = os.path.join(checkpoint_dir, "model")
        model = tf.keras.models.load_model(checkpoint_filepath)
    else:
        model = tf.keras.models.Sequential([
            tf.keras.layers.Flatten(input_shape=(28, 28)),
            tf.keras.layers.Dense(config["hidden"], activation="relu"),
            tf.keras.layers.Dropout(0.2),
            tf.keras.layers.Dense(num_classes, activation="softmax")
        ])

        model.compile(
            loss="sparse_categorical_crossentropy",
            optimizer=tf.keras.optimizers.SGD(
                lr=config["lr"], momentum=config["momentum"]),
            metrics=["accuracy"])
    return model

それでは用意したトレーニングループを使ってPopulation Based Trainingを実行しましょう。複数プロセスを管理するコントローラを書くのは骨が折れる作業ですが、Ray Tuneならへっちゃらです。
tune.runの引数に①さっき作ったトレーニングループの関数と②おまじないで作るPopulationBasedTrainingインスタンスを渡してやるだけ！

sched = PopulationBasedTraining(
    time_attr="training_iteration",
    perturbation_interval=4,
    hyperparam_mutations=search_space  # 後述
)

analysis = tune.run(
    train_mnist,
    name="exp",
    scheduler=sched,
    metric="mean_accuracy",
    mode="max",
    stop={
        "training_iteration": 10  # エージェント本番はもうちょっと増やしましょう。
    },
    num_samples=10,  # 並走するエージェントの数。大きくしても同時に保持するプロセスの数は変わらない。
    resources_per_trial=resource_per_trial,  # 後述
)
print("Best hyperparameters found were: ", analysis.best_config)

説明のない引数についてはおまじないだとおもっていてください。

探索空間の設定

PopulationBasedTrainingの引数にあったhyperparam_mutations引数にいい感じのdictを渡してあげることでハイパラの探索空間が指定できます。連続的なものから離散的なもの、はては自作関数まで幅広く使えるのですが詳しくはドキュメントを参照。

search_space = {
    "lr": tune.uniform(0.001, 0.1),
    "momentum": tune.uniform(0.1, 0.9),
    "hidden": tune.randint(32, 512),
}

計算コストの割当

tune.runのresources_per_trial引数にわたすdictで各プロセスあたりの計算資源を設定できるのですが...ここは少しだけ気を使っていただきたい。

例えば８コア１GPUのマシンなら

resource_per_trial = {
    "cpu": 2,
    "gpu": 0.25
}

こんな設定はとても無難です。上のサンプルなら皆さんのマシンでそのまま動かすことができるでしょう。

しかしながらあなたの学習をこのリソース配分で実行して本当に大丈夫でしょうか？

自作トレーニングループのGPU使用メモリ量が4GBだとしたら、GPU１枚には16GBが必要になりますね。

GPUメモリが6GBしかないのにそんな無茶をさせてしまうとあっさりとフリーズすることもあります。

うっかり会社で借りてるクラウドをダウンさせてしまわないように注意してくださいね。

まずは学習ループ１つを走らせてみてnvidia-smiしてGPU使用量を確認してから適切な数値を設定してあげてください。小さすぎたらダメですよ。

まとめ

筆者のうろおぼえな記憶ではこのチュートリアルコードでかいたような単純なMLPではMNISTは97％くらいしかでなかったようなきがします。
でもハイパラを本気でチューニングしてみると９９％くらい出ます。

Ray Tuneは使ってみたらとても簡単に実装できたのでハイパラ最適化に困っている人はぜひ使ってみてください。

あとはTensorflow2以外に応用する場合などに向けての細かいチュートリアルは書きませんでしたが需要があればかくかもしれません。

それでは。