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

このガイドでは、ハイパーパラメーター探索を最適化するために、W&B を Python のトレーニングスクリプトやノートブックに統合する際の推奨事項を紹介します。

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

モデルをトレーニングする Python スクリプトがあるとします (以下を参照) 。ここでの目的は、検証精度 (val_acc) を最大化するハイパーパラメーターを検索することです。 Python スクリプトでは、train_one_epoch と evaluate_one_epoch の 2 つの関数を定義します。train_one_epoch 関数は 1 エポック分のトレーニングをシミュレートし、トレーニング精度と損失を返します。evaluate_one_epoch 関数は検証データセットでモデルを評価する処理をシミュレートし、検証精度と損失を返します。また、学習率 (lr) 、バッチサイズ (batch_size) 、エポック数 (epochs) などのハイパーパラメーター値を含む設定辞書 (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)
        print("トレーニング精度:", train_acc, "トレーニング損失:", train_loss)
        print("検証精度:", val_acc, "検証損失:", val_loss)

if __name__ == "__main__":
    main()

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

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

トレーニングスクリプトを更新して、W&B を組み込みます。W&B を Python スクリプトまたはノートブックにどのように統合するかは、sweeps の管理方法によって異なります。 sweeps の開始、停止、管理に W&B Python SDK を使用するには、Python script or notebook タブの手順に従ってください。代わりに W&B CLI を使用する場合は、CLI タブの手順に従ってください。

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

sweep の設定を記述した YAML 設定ファイルを作成します。この設定ファイルには、sweep で探索するハイパーパラメーターを記述します。次の例では、各 sweep の実行で、バッチサイズ (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]

W&B sweep 設定の作成方法の詳細については、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() でログします。設定で定義したメトリクスを必ずログする必要があります。この例の設定 dict (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 agent が試行する run の最大数を設定します。これは省略可能です。この例では、最大数を 5 に設定しています。

NUM=5

続いて、wandb sweep コマンドで sweep を初期化します。YAML ファイル名を指定します。必要に応じて、プロジェクトフラグ (--project) にプロジェクト名も指定します。

wandb sweep --project sweep-demo-cli config.yaml

これにより、sweep ID が返されます。sweep の初期化方法について詳しくは、 sweeps を初期化するを参照してください。sweep ID をコピーし、次のコードスニペット内の sweepID を置き換えて、 wandb agent コマンドで sweep ジョブを開始します:

wandb agent --count $NUM your-entity/sweep-demo-cli/sweepID

詳細は、sweep ジョブを開始するを参照してください。

次の step に従って、PythonスクリプトにW&Bを追加します。

キーと値のペアで sweep 設定を定義する辞書オブジェクトを作成します。sweep 設定では、W&B がユーザーに代わって探索するハイパーパラメーターと、最適化したいメトリクスを定義します。前の例に続いて、各 sweep で変化させるハイパーパラメーターは、バッチサイズ (batch_size) 、エポック数 (epochs) 、および学習率 (lr) です。検証スコアの精度を最大化したいので、"goal": "maximize" を設定し、最適化対象の変数名を指定します。この例では val_acc ("name": "val_acc") です。
wandb.sweep() に sweep の設定辞書を渡します。これにより sweep が初期化され、sweep ID (sweep_id) が返されます。詳細は、sweep の初期化を参照してください。
スクリプトの先頭で、W&B Python SDK (wandb) をimportします。
main 関数内で、wandb.init() を使用して、データを W&B Run として Sync およびログするバックグラウンドプロセスを生成します。プロジェクト名をパラメーターとして wandb.init() method に渡します。プロジェクト名を渡さない場合、W&B はデフォルトのプロジェクト名を使用します。
wandb.Run.config オブジェクトからハイパーパラメーターの値を取得します。これにより、ハードコードした値ではなく、sweep の設定 dict で定義したハイパーパラメーターの値を使用できます。
最適化対象のメトリクスを、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 agentから呼び出された場合、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プロジェクト名")
    args = parser.parse_args()

    # sweep設定のdictを定義する
    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},
        },
    }

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

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

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

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

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