PyTorch

Try in Colab

PyTorch は、特に研究者の間で、Python におけるディープラーニングの最も人気のあるフレームワークの一つです。W&B は、PyTorch に対して一流のサポートを提供し、勾配のログから CPU と GPU 上でのコードのプロファイリングまで対応しています。

Colab ノートブックで私たちのインテグレーションを試してみてください。

Try in Colab

また、example repo では、スクリプトや Fashion MNIST を使用した Hyperband によるハイパーパラメータ最適化などの例を含むものがあります。それが生成する W&B Dashboard もご覧いただけます。

wandb.watch を使った勾配のログ

勾配を自動的にログするには、wandb.watch を呼び出して、PyTorch モデルを渡します。

import wandb

wandb.init(config=args)

model = ...  # モデルをセットアップする

# マジック
wandb.watch(model, log_freq=100)

model.train()
for batch_idx, (data, target) in enumerate(train_loader):
    output = model(data)
    loss = F.nll_loss(output, target)
    loss.backward()
    optimizer.step()
    if batch_idx % args.log_interval == 0:
        wandb.log({"loss": loss})

同じスクリプト内で複数のモデルを追跡する必要がある場合は、それぞれのモデルに対して wandb.watch を個別に呼び出すことができます。この関数の参照ドキュメントはこちら。

勾配、メトリクス、およびグラフは、フォワード および バックワードパスの後に wandb.log が呼び出されるまでログされません。

画像とメディアのログ

画像データを持つ PyTorch Tensors を wandb.Image に渡すことができ、torchvision のユーティリティが自動的に画像に変換します。

images_t = ...  # PyTorch Tensors として画像を生成またはロードする
wandb.log({"examples": [wandb.Image(im) for im in images_t]})

PyTorch や他のフレームワークにおけるリッチメディアのログについての詳細は、メディアログガイドをご覧ください。

メディアと一緒にモデルの予測や派生メトリクスなどの情報も含めたい場合は、wandb.Table を使用します。

my_table = wandb.Table()

my_table.add_column("image", images_t)
my_table.add_column("label", labels)
my_table.add_column("class_prediction", predictions_t)

# Table を W&B にログ
wandb.log({"mnist_predictions": my_table})

データセットやモデルのログと視覚化についての詳細は、W&B Tables のガイドをご覧ください。

PyTorch コードのプロファイリング

W&B は PyTorch Kineto の Tensorboard プラグイン と直接統合されており、PyTorch コードのプロファイリング、CPU と GPU の通信の詳細の検査、ボトルネックや最適化を識別するためのツールを提供します。

profile_dir = "path/to/run/tbprofile/"
profiler = torch.profiler.profile(
    schedule=schedule,  # スケジュールの詳細はプロファイラードキュメントを参照
    on_trace_ready=torch.profiler.tensorboard_trace_handler(profile_dir),
    with_stack=True,
)

with profiler:
    ...  # プロファイルしたいコードをここで実行
    # 詳細な使用情報はプロファイラードキュメントを参照

# wandb アーティファクトを作成
profile_art = wandb.Artifact("trace", type="profile")
# pt.trace.json ファイルをアーティファクトに追加
profile_art.add_file(glob.glob(profile_dir + ".pt.trace.json"))
# アーティファクトをログ
profile_art.save()

こちらの Colabで作業中の例コードを見て実行できます。

インタラクティブなトレースビューツールは、Chrome Trace Viewer に基づいており、Chrome ブラウザで最も良好に動作します。