Ray Tune でハイパーパラメータを自動調整する

Ray は Apache-2.0 license *1 で公開されている， AI 関連や Python の処理を簡単に並列化して大規模に実行可能にするフレームワークです*2．この中に Ray Tune というライブラリがあります．

Ray Tune: Hyperparameter Tuning
Ray Tune: Hyperparameter Tuning — Ray 2.7.1

Ray は複数のライブラリで構成されており，単独でインストールして使用できます．

$ pip install "ray[tune]" torch

tune.Tunerが基礎となるクラスで，ハイパーパラメータを自動調整する時に評価する指標 (trainable) ，ハイパーパラメータを指定する辞書 (param_space) ，実行時の設定 (train.RunConfig)，自動調整の設定 (tune.TuneConfig) を指定し， tuner.fit() 関数を実行することで，trainable 関数が最小/最大となる条件を見つけてくれます．

環境

この記事では以下のバージョンのライブラリを使用しました． Ray は現在 (2023/10/18) 非常に活発に開発されているため，将来のバージョンでは動作が変更になる可能性があります*3．

Python 3.11.6 for MacOS 13.6
torch==2.1.0
ray==2.7.1

簡単な関数の最小値を求める

a × 5 + b が最小値となる条件を求めてみます．まずは trainable 関数を定義します．

from ray import train, tune

def func(config):
    return {
        "score": config["a"] * 5 + config["b"],
    }

これが関数ですので，自動調整する処理を作成します．

tuner = tune.Tuner(
    func,
    param_space={
        "a": tune.grid_search([0, 1, 2, 3]),
        "b": tune.grid_search([0, 5, 10]),
    },
    run_config=train.RunConfig(
        name="test01",
    ),
    tune_config=tune.TuneConfig(
        mode="min",
        max_concurrent_trials=1,
    ),
)
results = tuner.fit()
print(results.get_best_result(metric="score", mode="min").config)

以下のようなログとともに結果が表示されました．

╭────────────────────────────────────────────────────────╮
│ Configuration for experiment     test01                │
├────────────────────────────────────────────────────────┤
│ Search algorithm                 BasicVariantGenerator │
│ Scheduler                        FIFOScheduler         │
│ Number of trials                 12                    │
╰────────────────────────────────────────────────────────╯

...

Trial status: 12 TERMINATED
Current time: 2023-10-18 21:50:26. Total running time: 1s
Logical resource usage: 1.0/8 CPUs, 0/0 GPUs
╭─────────────────────────────────────────────────────────────────────────────────╮
│ Trial name         status         a     b     iter     total time (s)     score │
├─────────────────────────────────────────────────────────────────────────────────┤
│ func_e7573_00000   TERMINATED     0     0        1        0.000130177         0 │
│ func_e7573_00001   TERMINATED     1     0        1        9.20296e-05         5 │
│ func_e7573_00002   TERMINATED     2     0        1        0.000113964        10 │
│ func_e7573_00003   TERMINATED     3     0        1        0.00020504         15 │
│ func_e7573_00004   TERMINATED     0     5        1        8.89301e-05         5 │
│ func_e7573_00005   TERMINATED     1     5        1        8.70228e-05        10 │
│ func_e7573_00006   TERMINATED     2     5        1        0.000291109        15 │
│ func_e7573_00007   TERMINATED     3     5        1        7.98702e-05        20 │
│ func_e7573_00008   TERMINATED     0    10        1        8.82149e-05        10 │
│ func_e7573_00009   TERMINATED     1    10        1        8.39233e-05        15 │
│ func_e7573_00010   TERMINATED     2    10        1        7.00951e-05        20 │
│ func_e7573_00011   TERMINATED     3    10        1        8.4877e-05         25 │
╰─────────────────────────────────────────────────────────────────────────────────╯

{'a': 0, 'b': 0}

a の候補を 3 種類，b の候補を 4 種類与えたので，12 個の条件について値を求めてくれました．結果として，最小となる条件が {'a': 0, 'b': 0} であることが分かりました．

tune.search.basic_variant.BasicVariantGenerator は，全ての組み合わせを実行して検証してくれる探索アルゴリズムです．様々なアルゴリズムが用意されていますので，お好みのものを使用します．

Tune Search Algorithms (tune.search) — Ray 2.7.1

ray.tune.schedulers.FIFOScheduler *4 は単純に作成された試行 (trial) を順番に実行します．

Tune Trial Schedulers (tune.schedulers) — Ray 2.7.1

param_space には，Tune Search Space API を使用した辞書を指定します．tune.grid_search を使えば指定した値の組み合わせを評価してくれます．これに対して，tune.loguniform(1e-4, 1e-2) ならば対数で 0.0001 から 0.01 までの数値を選択して評価してくれますし，tune.choice ならば指定した値の中から選択して評価してくれます．選択する数は，tune.TuneConfig の num_samples *5 で指定します．

Tune Search Space API — Ray 2.7.1

train.report, train.Checkpoint

Checkpoint はユーザーが定義する訓練途中のスナップショットです．適当なディレクトリにモデルのデータを出力して，場所を train.report() 関数で指定することで， Ray Tune が管理してくれるようになります．

先程の例を以下のように置き換えます．

def func(config):
    metrics = {
        "score": config["a"] * 5 + config["b"],
    }
    for _ in range(5):
        with tempfile.TemporaryDirectory() as tempdir:
            torch.save(
               {},
                os.path.join(tempdir, "checkpoint.pt"),
            )
            train.report(
                metrics=metrics,
                checkpoint=train.Checkpoint.from_directory(tempdir),
            )
    return metrics

分散処理するときに同じディレクトリに上書きしないように，tempfile モジュールを利用して，専用のディレクトリに保存するのがよいです．自動的に，{storage_path}/{name}/ に移動されるようです*6．

保存した Checkpoint は読み込めます．先程保存した checkpoint.pt は以下のようにして読み込めます．

def func(config):
    checkpoint = train.get_checkpoint()
    if checkpoint:
        with checkpoint.as_directory() as checkpoint_dir:
            checkpoint_dict = torch.load(os.path.join(checkpoint_dir, "checkpoint.pt"))

Ray Tune 2.7 の変更

以前は Ray AIR session という機能を利用してスナップショットを記録する仕組みだったようですが， Ray Tune 2.7 から train, tune などに再配置されたようです*7．2.7 以前のサンプルとは書き方が異なるので注意が必要です．

結果を可視化する

Ray Tune には可視化ツールは含まれていないようですが，TensorBoard に対応してるようです*8．

tensorboard --logdir {PATH}

このコマンドで以下のように可視化できます．便利ですね．

GPU を割り当てる

tune.with_resources(trainable, resources_per_trial) で gpu を指定すると，環境変数 CUDA_VISIBLE_DEVICES がいい感じに設定されます*9．PyTorch は当該環境変数を認識してくれます*10．あとは PyTorch のやり方に従い， model.to() や Tensor.to() を使用して GPU にコピーすればよいです．