コードに W&B (wandb) を追加する - Weights & Biases Documentation

このガイドでは、ハイパーパラメーター探索を最適化するために、W&B を Python のトレーニングスクリプトやノートブックに統合する際の推奨事項を紹介します。これらの推奨事項に従うことで、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 = {"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 を追加する方法を説明します。ここでは、検証精度 (val_acc) を最大化する最適なハイパーパラメーターを見つけるために W&B を使用します。

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

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

CLI
Python スクリプトまたはノートブック

sweep の設定を含む YAML 設定ファイルを作成します。この設定ファイルには、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) を import します。PyYAML を使用して YAML 設定ファイルを読み込みます。
設定ファイルを読み込みます。
wandb.init() を使用して、データの sync とログを行うバックグラウンドプロセスを開始し、W&B Run を作成します。config オブジェクトを config パラメーターに渡します。
ハイパーパラメーターの値は、ハードコードされた値を使用する代わりに wandb.Run.config から取得します。
wandb.Run.log() を使用して、最適化したいメトリクスをログします。設定で定義したメトリクスを必ずログする必要があります。この例の設定 dictionary (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 が返されます。詳細は、Initialize sweeps を参照してください。
sweep ID をコピーし、次のコマンドの [SWEEP-ID] を置き換えて、wandb agent コマンドで sweep ジョブを開始します。[YOUR-ENTITY] は W&B のエンティティ名に置き換えます。必要に応じて、--count を設定して、エージェントが試行する Runs の数を制限します (この例では 5 に設定しています) 。
wandb agent --count 5 [YOUR-ENTITY]/sweep-demo-cli/[SWEEP-ID]

sweep エージェントはトレーニングスクリプトを繰り返し実行し、そのたびに YAML 設定内の異なるハイパーパラメーター値の組み合わせを使用して、結果を W&B にログします。詳細は、Start sweep jobs を参照してください。

次の手順に従って、W&B を Python スクリプトに追加してください。

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

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 を追加する