Python ライブラリに W&B を追加する - Weights & Biases Documentation

このガイドでは、ユーザーがあなたのコードを使用する際に、実験をトラッキングし、システムメトリクスを監視し、モデルを管理できるように、Python ライブラリに W&B を統合する方法を説明します。対象読者は、独自のフレームワーク、SDK、または再利用可能なトレーニングコードを通じて W&B の機能を利用できるようにしたい、ライブラリの作成者やメンテナーです。単一の Python トレーニングスクリプトや Jupyter ノートブックよりも複雑なコードベース (トレーニングフレームワーク、SDK、または再利用可能なライブラリなど) に W&B を統合する場合は、以下の推奨事項に従ってください。

W&B を初めて使用する場合は、先にコアガイド (たとえば Experiment Tracking) を確認してください。

以下のセクションでは、主要なインテグレーション上の判断ポイントを順に説明します。具体的には、W&B のインストール方法、認証方法、run の開始方法と設定方法、メトリクスやアーティファクトをログする方法、そして分散トレーニングやハイパーパラメーター sweep をどのようにサポートするかを扱います。

ユーザーによるW&Bのインストール方法を決める

開始する前に、W&Bをライブラリの必須依存関係にするか、オプション機能として扱うかを決めてください。この選択は、W&B Python SDK (wandb) をどのようにimportするか、インストール方法をどのようにドキュメント化するか、また wandb が存在しない環境をどのように処理するかに影響します。

W&B を依存関係として必須にする

W&B がライブラリの中核機能に関わる場合は、ライブラリと一緒に自動的にインストールされるよう、依存関係に wandb を追加します。

torch==1.8.0 
...
wandb==0.13.*

インストール時に W&B をオプションにする

W&B がオプション機能である場合は、実験管理を必要としない Users も引き続きコードを使用できるように、インストールされていなくてもライブラリが動作するようにします。 Python で wandb を条件付きで import するか、pyproject.toml でオプションの依存関係として宣言できます。

Python
pyproject.toml

wandb が利用可能かどうかを確認し、インストールしていない状態でユーザーが W&B の機能を有効にした場合は、明確なエラーを出します。

try:
    import wandb
    _WANDB_AVAILABLE = True
except ImportError:
    _WANDB_AVAILABLE = False

pyproject.toml ファイルで wandb をオプションの依存関係として宣言します。

[project]
name = "my_awesome_lib"
version = "0.1.0"
dependencies = [
    "torch",
    "sklearn"
]

[project.optional-dependencies]
dev = [
    "wandb"
]

ユーザー認証

W&B では、APIキーを使用してユーザーやマシンを認証します。ユーザーがライブラリから Runs をログするには、事前に APIキーを発行し、wandb クライアントで使用できるようにしておく必要があります。

APIキーを発行する

APIキーは、クライアントやマシンをW&Bに認証するために使用します。続いて行うログイン手順で使用できるように、ユーザープロフィールからAPIキーを発行します。

より手早く行うには、User Settings にアクセスしてAPIキーを作成してください。APIキーはすぐにコピーし、パスワードマネージャーなどの安全な場所に保存してください。

右上にあるユーザープロフィールアイコンをクリックします。
User Settings を選択し、API Keys セクションまでスクロールします。

W&B をインストールしてログインする

APIキーを取得したら、wandb ライブラリをローカルにインストールしてログインし、以降の Runs が W&B に対して認証できるようにします。ご利用の環境に合ったタブを選択してください。

コマンドライン
Python
Python notebook

WANDB_API_KEY の環境変数に APIキーを設定します。<> で囲まれた値はご自身の値に置き換えてください。
export WANDB_API_KEY=<your_api_key>
wandb ライブラリをインストールしてログインします。
pip install wandb wandb login

ターミナルで Python SDK をインストールします。
pip install wandb
Python スクリプトまたはノートブックから W&B にログインします。APIキーの入力を求められます。
import wandb wandb.login()

次のコードスニペットを Jupyter notebook のセルにコピー＆ペーストして実行します。APIキーの入力を求められます。

!pip install wandb

import wandb
wandb.login()

run を開始する

認証を設定したら、次のステップは W&B run を開始することです。これにより、ライブラリがメトリクス、設定、アーティファクトをログする場所を確保できます。 run は、トレーニング実験などの単一の計算単位を表します。ほとんどのライブラリでは、トレーニングジョブごとに 1 つの run が作成されます。run の詳細については、W&B Runs を参照してください。 wandb.init() で run を初期化し、プロジェクト名とチーム entity (チーム名) を指定します。プロジェクトを指定しない場合、W&B は run を “uncategorized” というデフォルトのプロジェクトに保存します。<> で囲まれた値は、ご自身の値に置き換えてください。

with wandb.init(project="<project_name>", entity="<entity>") as run:
    ...

W&B では、エラーが発生した場合でも run が確実に適切に終了されるよう、コンテキストマネージャーを使用することを推奨しています。コンテキストマネージャーを使用しない場合は、run.finish() を呼び出して run を終了し、すべてのデータを W&B にログする必要があります。run を終了すると、プロセスが終了する前に、すべてのメトリクス、設定、アーティファクトがアップロードされることが保証されます。

wandb.init() を呼び出すタイミングwandb.init() はできるだけ早く呼び出してください。W&B は stdout、stderr、エラーメッセージをキャプチャするため、デバッグが容易になります。関連するすべての情報が run にキャプチャされるよう、トレーニングループ全体を wandb.init() のコンテキストマネージャーで囲んでください。これにはエラーメッセージも含まれており、デバッグで重要になることがあります。

`wandb` をオプションの依存関係にする

実行時に wandb をオプションにして、ユーザーが W&B runs を生成せずにライブラリを実行できるようにしたい場合は、次のいずれかの方法を使用できます。

wandb フラグを定義します。
wandb.init() で wandb を disabled に設定します。
wandb をオフラインに設定します。なお、この場合も wandb 自体は実行されますが、インターネット経由で W&B と通信しません。

たとえば、次のように wandb フラグを定義します。

Python
Bash

trainer = my_trainer(..., use_wandb=True)

python train.py ... --use-wandb

wandb.init() で wandb を disabled に設定します。

Python
Bash

wandb.init(mode="disabled")

export WANDB_MODE=disabled

または

wandb disabled

wandb をオフラインに設定します。

環境変数
Bash

export WANDB_MODE=offline

または

os.environ['WANDB_MODE'] = 'offline'

wandb offline

run 設定を定義する

run を初期化した後、その run に関連付けられたハイパーパラメーターやその他のメタデータを記録する設定辞書を追加できます。設定をログすると、後から run を比較、フィルター、再現しやすくなります。 run の初期化時に設定辞書を指定すると、ハイパーパラメーターやその他のメタデータを W&B にログできます。 W&B では、設定パラメーターに基づいて run を比較したり、Runs table でフィルターしたりできます。また、これらのパラメーターを使って、W&B で run をグループ化することもできます。たとえば、次の画像では、バッチサイズ (batch_size) が設定パラメーターとして定義されており、Runs table に表示されています (最初の列を参照) 。これにより、ユーザーはバッチサイズに基づいて run をフィルターし、比較できます。

一般的な設定パラメーターの値には、次のようなものがあります。

モデル名、バージョン、アーキテクチャーパラメーター、ハイパーパラメーター。
データセット名、バージョン、トレーニングまたは検証のサンプル数。
学習率、バッチサイズ、オプティマイザーなどのトレーニングパラメーター。

次のコードスニペットは、設定をログする方法を示しています。

config = {"batch_size": 32, ...}
with wandb.init(..., config=config) as run:
    ...

run 設定を更新する

モデルのパラメーター数など、一部の設定値は wandb.init() を呼び出す時点ではまだわからないことがあります。初期化時に値を利用できない場合は、後から wandb.Run.config.update を使って設定を更新します。たとえば、モデルをインスタンス化した後で、そのパラメーターを追加したいことがあります。

with wandb.init(...) as run:
    model = MyModel(...)
    run.config.update({"model_parameters": 3500})

詳細は、Experiments の設定を参照してください。

メトリクスとデータをログする

run を開始して設定したら、メトリクスやその他のデータをログできるようになります。これにより、W&B はそれらを run に紐付けて記録します。

メトリクスをログする

損失や accuracy などのスカラーメトリクスをログするには、キーがメトリクスの名になる辞書を作成します。この辞書オブジェクトを wandb.Run.log() に渡して、W&B にログします。

NUM_EPOCHS = 10

for epoch in range(NUM_EPOCHS):
    for input, ground_truth in data: 
        prediction = model(input) 
        loss = loss_fn(prediction, ground_truth) 
        metrics = { "loss": loss } 
        run.log(metrics)

メトリクス名に接頭辞を付けると、W&B で関連するメトリクスをグループ化できます。一般的な接頭辞としては、トレーニング用メトリクスの train/ や検証用メトリクスの val/ などがありますが、ユースケースに応じて任意の接頭辞を使用できます。これにより、project のWorkspaceに、トレーニング用メトリクスと検証用メトリクス、または分けて表示したい他のタイプのメトリクスごとに、個別のセクションが作成されます:

with wandb.init(...) as run:
    metrics = {
        "train/loss": 0.4,
        "train/learning_rate": 0.4,
        "val/loss": 0.5, 
        "val/accuracy": 0.7
    }
    run.log(metrics)

詳しくは wandb.Run.log() を参照してください。

x-axis を制御する

デフォルトでは、W&B Python SDK は独自の step カウンタを管理しており、これはトレーニングループにおける step の意味と一致しない場合があります。同じトレーニング step に対して wandb.Run.log() を複数回呼び出すと、wandb SDK は wandb.Run.log() を呼び出すたびに内部の step カウンタをインクリメントします。このカウンタは、トレーニングループ内のトレーニング step と一致しないことがあります。これを避けるには、wandb.init() を呼び出した直後に、wandb.Run.define_metric() を使って x-axis の step を一度だけ明示的に定義します:

with wandb.init(...) as run:
    run.define_metric("*", step_metric="global_step")

グロブパターン * は、すべてのメトリクスでグラフの x 軸に global_step が使用されることを意味します。global_step に対してログするメトリクスを特定のものだけにしたい場合は、代わりにそれらを指定できます。

run.define_metric("train/loss", step_metric="global_step")

次に、wandb.Run.log() を呼び出すたびに、メトリクス、step メトリクス、および global_step をログしてください:

for step, (input, ground_truth) in enumerate(data):
    ...
    run.log({"global_step": step, "train/loss": 0.1})
    run.log({"global_step": step, "eval/loss": 0.2})

独立したstep変数にアクセスできない場合、たとえば検証ループ中に global_step を使用できないときは、以前にログした global_step の値が W&B によって自動的に使用されます。この場合は、必要になったときに定義済みになっているよう、あらかじめメトリクスの初期値をログしておいてください。

メディアと構造化データをログする

スカラーに加えて、画像、表、テキスト、オーディオ、動画などもログできます。メトリクスとあわせてメディアをログすることで、ユーザーは時間の経過に伴うモデルの定性的な動作を確認できます。データをログする際の留意点には、次のようなものがあります。

メトリクスはどのくらいの頻度でログするべきですか。任意にするべきですか。
可視化に役立つのは、どのようなタイプのデータですか。
- 画像については、サンプル予測やセグメンテーションマスクをログすることで、時間の経過に伴う変化を確認できます。
- テキストについては、後で詳しく調べられるように、サンプル予測の表をログできます。

詳細は、オブジェクトとメディアをログするを参照してください。

分散トレーニングをサポートする

ライブラリが複数のプロセスまたはマシンにまたがってトレーニングを実行できる場合は、その環境でW&Bをどのように動作させるかを決めて、ログの整合性を保ち、重複しないようにしてください。分散環境をサポートするフレームワークでは、次のいずれかのワークフローを利用できます。

メインプロセスからのみログする (推奨) 。
すべてのプロセスからログし、共有のgroup名でrunをグループ化する。

詳細は分散トレーニングのExperimentsをログするを参照してください。

アーティファクトでモデルとデータセットをトラッキングする

メトリクスに加えて、ユーザーが Runs を再現して比較できるように、ライブラリが生成または使用するモデルとデータセットを永続化できます。 W&B Artifacts を使用して、モデルとデータセットをトラッキングし、バージョン管理できます。アーティファクトは機械学習アセットのストレージとバージョン管理を提供し、データとモデルの関連性を示すリネージを自動的にトラッキングします。

ライブラリにアーティファクトを統合する際は、次の点を考慮してください。

モデルチェックポイントまたはデータセットをアーティファクトとしてログするかどうか (オプションにする場合) 。
Artifact の入力参照 (たとえば、entity/project/artifact) 。
モデルチェックポイントまたはデータセットをログする頻度。たとえば、各エポックまたは 500 step ごとです。

モデルチェックポイントをログする

モデルチェックポイントをアーティファクトとしてログすると、ユーザーはトレーニング済みの重みを復元、バージョン管理、共有できます。一般的には、W&B が生成する一意の run ID をアーティファクト名の一部として使用し、チェックポイントをアーティファクトとしてログします。

metadata = {"eval/accuracy": 0.8, "train/steps": 800} 

artifact = wandb.Artifact(
                name=f"model-{run.id}", 
                metadata=metadata, 
                type="model"
                ) 
artifact.add_dir("output_model") # モデルの重みが保存されているローカルディレクトリ

aliases = ["best", "epoch_10"] 
run.log_artifact(artifact, aliases=aliases)

前のスニペットでは、モデルチェックポイントをアーティファクトとしてログし、評価精度やトレーニング step などのメタデータを追加します。このアーティファクトの名には一意のrun IDが含まれ、すばやく参照できるようにカスタムエイリアスが設定されます。

入力アーティファクトをログする

データとモデルのリネージを取得するには、run が入力として使用するデータセットまたは事前トレーニング済みモデルをログします:

dataset = wandb.Artifact(name="flowers", type="dataset")
dataset.add_file("flowers.npy")
run.use_artifact(dataset)

前のスニペットでは、“flowers” という名前のデータセット用のアーティファクトを作成し、そのアーティファクトにファイルを追加しています。wandb.Run.use_artifact() の呼び出しによって、そのアーティファクトを現在の run に関連付けます。これにより、W&B はその run で使用されたデータセットのリネージをトラッキングできます。

アーティファクトをダウンロードする

アーティファクトをログした後、以前にW&Bにログしたアーティファクトをダウンロードし、トレーニングまたは推論のコードで使用できます。適切な方法は、すでにアクティブなrunがあるかどうかによって異なります。 runコンテキストがある場合は、wandb.Run.use_artifact() を使用してW&B内のアーティファクトを参照し、その後 wandb.Artifact.download() を呼び出してローカルディレクトリにダウンロードします。wandb.Run.use_artifact() を使用すると、artifact は現在のrunへの入力としても記録され、リネージが保持されます。

with wandb.init(...) as run:
    artifact = run.use_artifact("user/project/artifact:latest")
    local_path = artifact.download()

W&B Public API を使用すると、runを初期化せずにアーティファクトを参照およびダウンロードできます。これは、分散環境や推論を行う場合など、新しいrunを作成したくないシナリオで役立ちます。

import wandb
artifact = wandb.Api().artifact("user/project/artifact:latest")
local_path = artifact.download()

詳細は、アーティファクトをダウンロードして使用するを参照してください。

ハイパーパラメーターを調整する

お使いのライブラリがハイパーパラメーターのチューニングをサポートしている場合は、W&B Sweepsを統合して、Experiments を管理・可視化できます。Sweeps は、定義された探索空間内で複数の run を連携して実行し、その結果を W&B 上に表示することで、ユーザーが設定を並べて比較できるようにします。

​ユーザーによるW&Bのインストール方法を決める

​W&B を依存関係として必須にする

​インストール時に W&B をオプションにする

​ユーザー認証

​APIキーを発行する

​W&B をインストールしてログインする

​run を開始する

​wandb をオプションの依存関係にする

​run 設定 を定義する

​run 設定 を更新する

​メトリクスとデータをログする

​メトリクスをログする

​x-axis を制御する

​メディアと構造化データをログする

​分散トレーニングをサポートする

​アーティファクト でモデルとデータセットをトラッキングする

​モデル チェックポイントをログする

​入力アーティファクトをログする

​アーティファクト をダウンロードする

​ハイパーパラメーターを調整する

ユーザーによるW&Bのインストール方法を決める

W&B を依存関係として必須にする

インストール時に W&B をオプションにする

ユーザー認証

APIキーを発行する

W&B をインストールしてログインする

run を開始する

`wandb` をオプションの依存関係にする

run 設定を定義する

run 設定を更新する

メトリクスとデータをログする

メトリクスをログする

x-axis を制御する

メディアと構造化データをログする

分散トレーニングをサポートする

アーティファクトでモデルとデータセットをトラッキングする

モデルチェックポイントをログする

入力アーティファクトをログする

アーティファクトをダウンロードする

ハイパーパラメーターを調整する