コードに W&B (wandb) を追加する

このガイドでは、ハイパーパラメーター探索を最適化するために、Python のトレーニングスクリプトまたはノートブックに W&B を統合する際の推奨事項を紹介します。これらの推奨事項に従うと、W&B Sweeps を使用してハイパーパラメーターの値を探索し、トレーニングと検証のメトリクスをログして、優れたモデル性能が得られる設定を特定できます。このガイドは、Python のトレーニングスクリプトをすでに用意しており、ハイパーパラメーター sweep のサポートを追加したい機械学習の実務者を対象としています。以下のセクションでは、まずトレーニングスクリプトの例を紹介し、次にそのスクリプトを W&B Sweeps で動作するように更新する方法を説明します。

元のトレーニングスクリプト

モデルをトレーニングする Python スクリプトがあるとします (以下のコードを参照) 。目標は、検証精度 (val_acc) を最大化するハイパーパラメーターを見つけることです。この Python スクリプトでは、train_one_epoch と evaluate_one_epoch の 2 つの関数を定義します。train_one_epoch 関数は 1 エポック分のトレーニングをシミュレートし、トレーニングの精度と損失を返します。evaluate_one_epoch 関数は検証データセットでのモデルの評価をシミュレートし、検証精度と損失を返します。学習率、バッチサイズ、エポック数などのハイパーパラメーター値を含む、config という名前の設定辞書を定義します。この設定辞書の値によってトレーニングプロセスを制御します。次に、一般的なトレーニングループを模した main という関数を定義します。スクリプトはエポックごとに、トレーニングデータセットと検証データセットでの精度と損失を計算します。

このコードはモックのトレーニングスクリプトです。実際にはモデルをトレーニングせず、精度と損失の値をランダムに生成してトレーニングプロセスをシミュレートします。このコードの目的は、トレーニングスクリプトに W&B を統合する方法を示すことです。

import random
import numpy as np

def train_one_epoch(epoch, lr, batch_size):
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss

def evaluate_one_epoch(epoch):
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss

# ハイパーパラメーターの値を格納する config 変数
config = {"lr": 0.0001, "batch_size": 16, "epochs": 5}

def main():
    lr = config["lr"]
    batch_size = config["batch_size"]
    epochs = config["epochs"]

    for epoch in np.arange(1, epochs):
        train_acc, train_loss = train_one_epoch(epoch, lr, batch_size)
        val_acc, val_loss = evaluate_one_epoch(epoch)

        print("epoch: ", epoch)
        print("training accuracy:", train_acc, "training loss:", train_loss)
        print("validation accuracy:", val_acc, "validation loss:", val_loss)

if __name__ == "__main__":
    main()

以下のセクションでは、Python スクリプトに W&B を追加して、トレーニング中のハイパーパラメーターとメトリクスをトラッキングする方法を説明します。ここでは、W&B を使用して、検証精度 (val_acc) を最大化する最適なハイパーパラメーターを検索することを目指します。

トレーニングスクリプトに W&B を追加する

このセクションでは、元のトレーニングスクリプトを変更して、sweep エージェントが各 run にハイパーパラメーターの値を渡し、得られたメトリクスを W&B で記録できるようにする方法を説明します。Python スクリプトまたはノートブックに W&B を統合する方法は、sweep の管理方法によって異なります。 W&B Python SDK を使用して sweep の開始、停止、管理を行う場合は、Python script or notebook タブの手順に従ってください。W&B CLI を使用する場合は、CLI タブの手順に従ってください。

CLI
Python script or notebook

sweep 設定を記述した YAML 設定ファイルを作成します。設定ファイルには、sweep で探索するハイパーパラメーターを定義します。次の例では、sweep は run ごとに batch_size、epochs、lr の各ハイパーパラメーターの値を変化させます。

# config.yaml
program: train.py
method: random
name: sweep
metric:
  goal: maximize
  name: val_acc
parameters:
  batch_size:
    values: [16, 32, 64]
  lr:
    min: 0.0001
    max: 0.1
  epochs:
    values: [5, 10, 15]

詳細については、sweep 設定を定義するを参照してください。YAML ファイルの program キーには、Python スクリプトの名前を指定する必要があります。次に、コード例に以下を追加します。

W&B Python SDK (wandb) と PyYAML (yaml) をインポートします。PyYAML は YAML 設定ファイルの読み込みに使用します。
設定ファイルを読み込みます。
wandb.init() を使用して、データを同期し W&B Run としてログするバックグラウンドプロセスを開始します。設定オブジェクトを config パラメーターに渡します。
ハイパーパラメーターの値はハードコードせず、wandb.Run.config から取得して定義します。
最適化したいメトリクスを wandb.Run.log() でログします。設定で定義したメトリクスは必ずログしてください。この例では、設定辞書 (sweep_configuration) で、val_acc の値を最大化するように sweep を定義しています。

import wandb
import yaml
import random
import numpy as np


def train_one_epoch(epoch, lr, batch_size):
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss


def evaluate_one_epoch(epoch):
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss


def main():
    # デフォルトのハイパーパラメーターを設定します
    with open("./config.yaml") as file:
        config = yaml.load(file, Loader=yaml.FullLoader)

    with wandb.init(config=config) as run:
        for epoch in np.arange(1, run.config['epochs']):
            train_acc, train_loss = train_one_epoch(epoch, run.config['lr'], run.config['batch_size'])
            val_acc, val_loss = evaluate_one_epoch(epoch)
            run.log(
                {
                    "epoch": epoch,
                    "train_acc": train_acc,
                    "train_loss": train_loss,
                    "val_acc": val_acc,
                    "val_loss": val_loss,
                }
            )

# main 関数を呼び出します。
main()

トレーニングスクリプトを更新したら、CLI から sweep を初期化して開始します。

wandb sweep コマンドで sweep を初期化します。引数には YAML ファイル名を指定します。必要に応じて、--project フラグでプロジェクト名を指定します。
wandb sweep --project sweep-demo-cli config.yaml
このコマンドを実行すると sweep ID が返されます。詳細については、sweep を初期化するを参照してください。
sweep ID をコピーし、次のコマンドの [SWEEP-ID] を置き換えてから、wandb agent コマンドで sweep ジョブを開始します。[YOUR-ENTITY] は W&B の entity 名に置き換えてください。必要に応じて、--count を指定すると、エージェントが試行する run の数を制限できます (この例では 5 に設定しています) 。
wandb agent --count 5 [YOUR-ENTITY]/sweep-demo-cli/[SWEEP-ID]

sweep エージェントは、YAML 設定で定義したハイパーパラメーター値の組み合わせを毎回変えながらトレーニングスクリプトを繰り返し実行し、その結果を W&B にログします。詳細については、sweep ジョブを開始するを参照してください。

Python スクリプトに W&B を追加するには、次の手順を実行します。

キーと値のペアで sweep 設定を定義する辞書オブジェクトを作成します。sweep 設定では、W&B に探索させるハイパーパラメーターと、最適化するメトリクスを定義します。前の例に引き続き、各 sweep で batch_size、epochs、lr の各ハイパーパラメーターを変化させます。検証精度を最大化するには、メトリクスの goal を maximize に、name を val_acc に設定します。
sweep 設定の辞書を wandb.sweep() に渡します。これにより sweep が初期化され、sweep ID (sweep_id) が返されます。詳細については、Initialize sweeps を参照してください。
スクリプトの先頭で、W&B Python SDK (wandb) をインポートします。
main 関数内で wandb.init() を使用して、データを同期し W&B Run としてログするバックグラウンドプロセスを生成します。wandb.init() メソッドには、プロジェクト名をパラメーターとして渡します。プロジェクト名を渡さない場合、W&B はデフォルトのプロジェクト名を使用します。
wandb.Run.config オブジェクトからハイパーパラメーターの値を取得します。これにより、ハードコードされた値の代わりに、sweep 設定の辞書で定義したハイパーパラメーターの値を使用できます。
wandb.Run.log() を使用して、最適化対象のメトリクスを W&B にログします。設定で定義したメトリクスは必ずログしてください。たとえば、最適化するメトリクスを val_acc と定義した場合は、val_acc をログする必要があります。メトリクスをログしないと、W&B は最適化を実行できません。この例では、設定の辞書 (sweep_configuration) で、val_acc の値を最大化するように sweep を定義しています。
wandb.agent() で sweep を開始します。sweep ID と sweep が実行する関数名 (function=main) を指定し、run の最大数を 4 (count=4) に設定します。

これらをまとめると、スクリプトはたとえば次のようになります。

import wandb # W&B Python SDK をインポートします
import numpy as np
import random
import argparse

def train_one_epoch(epoch, lr, batch_size):
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss

def evaluate_one_epoch(epoch):
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss

def main(args=None):
    # sweep エージェントから呼び出された場合は args が None になるため、
    # sweep の設定で指定された project を使用します
    project = args.project if args else None
    
    with wandb.init(project=project) as run:
        # `wandb.Run.config` オブジェクトからハイパーパラメーターの値を取得します
        lr = run.config["lr"]
        batch_size = run.config["batch_size"]
        epochs = run.config["epochs"]

        # トレーニングループを実行し、パフォーマンスの値を W&B にログします
        for epoch in np.arange(1, epochs):
            train_acc, train_loss = train_one_epoch(epoch, lr, batch_size)
            val_acc, val_loss = evaluate_one_epoch(epoch)
            run.log(
                {
                    "epoch": epoch,
                    "train_acc": train_acc,
                    "train_loss": train_loss,
                    "val_acc": val_acc, # 最適化対象のメトリクス
                    "val_loss": val_loss,
                }
            )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--project", type=str, default="sweep-example", help="W&B project name")
    args = parser.parse_args()

    # sweep の設定の辞書を定義します
    sweep_configuration = {
        "method": "random",
        "name": "sweep",
        # 最適化するメトリクス
        # たとえば検証精度を最大化する場合は、
        # "goal": "maximize" と、最適化対象の変数名
        # (この例では "val_acc") を設定します
        "metric": {
            "goal": "maximize",
            "name": "val_acc"
            },
        "parameters": {
            "batch_size": {"values": [16, 32, 64]},
            "epochs": {"values": [5, 10, 15]},
            "lr": {"max": 0.1, "min": 0.0001},
        },
    }

    # 設定の辞書を渡して sweep を初期化します
    sweep_id = wandb.sweep(sweep=sweep_configuration, project=args.project)

    # sweep ジョブを開始します
    wandb.agent(sweep_id, function=main, count=4)

このスクリプトを実行すると、W&B が sweep を開始し、ハイパーパラメーターの組み合わせを変えながら main 関数を最大 4 回実行して、各 run のメトリクスをログします。ログされた結果は W&B App で比較できます。

sweep で W&B にメトリクスをログする最適化の対象として定義するメトリクスは、sweep 設定で定義するだけでなく、wandb.Run.log() でもログする必要があります。たとえば、sweep 設定で最適化するメトリクスを val_acc と定義した場合は、val_acc も W&B にログする必要があります。メトリクスをログしないと、W&B は最適化を実行できません。

with wandb.init() as run:
    val_loss, val_acc = train()
    run.log(
        {
            "val_loss": val_loss,
            "val_acc": val_acc
            }
        )

次の例は、W&B へのメトリクスのログ方法が誤っています。sweep 設定では val_acc を最適化の対象としていますが、コードでは val_acc をキー validation 配下のネストされた辞書内にログしています。メトリクスはネストされた辞書内ではなく、直接ログする必要があります。

with wandb.init() as run:
    val_loss, val_acc = train()
    run.log(
        {
            "validation": {
                "val_loss": val_loss, 
                "val_acc": val_acc
                }
            }
        )

​元のトレーニングスクリプト

​トレーニングスクリプトに W&B を追加する

元のトレーニングスクリプト

トレーニングスクリプトに W&B を追加する