PyTorch Lightning
6 minute read
PyTorch Lightning は、PyTorch コードを整理し、分散トレーニングや 16 ビット精度のような高度な機能を簡単に追加するための軽量ラッパーを提供します。 W&B は、あなたの ML 実験を記録するための軽量ラッパーを提供します。しかし、自分でそれらを組み合わせる必要はありません。Weights & Biases は、WandbLogger
を介して PyTorch Lightning ライブラリに直接組み込まれています。
Lightning との統合
wandb.log() を使用する際の注意点: WandbLogger
は Trainer の global_step
を使用して W&B にログを記録します。コード内で直接 wandb.log
を追加で呼び出す場合、wandb.log()
の step
引数を使用しないでください。
代わりに、Trainer の global_step
を他のメトリクスと同様に記録してください:

サインアップして APIキーを作成する
APIキー は、あなたのマシンを W&B に認証するためのものです。あなたのユーザープロフィールから APIキー を生成できます。
- 右上のユーザープロフィールアイコンをクリックします。
- User Settings を選択し、API Keys セクションまでスクロールします。
- Reveal をクリックします。表示された APIキー をコピーします。APIキー を非表示にするには、ページをリロードします。
wandb
ライブラリをインストールしてログインする
ローカルに wandb
ライブラリをインストールしてログインする方法:
PyTorch Lightning の WandbLogger
を使用する
PyTorch Lightning には、メトリクスやモデルの重み、メディアなどを記録するための複数の WandbLogger
クラスがあります。
Lightning と統合するには、WandbLogger をインスタンス化し、Lightning の Trainer
または Fabric
に渡してください。
よく使用されるロガーの引数
以下に、WandbLogger でよく使用されるパラメータを示します。すべてのロガー引数の詳細については PyTorch Lightning のドキュメントを確認してください。
Parameter | Description |
---|---|
project |
記録する wandb Project を定義します |
name |
あなたの wandb run に名前を付けます |
log_model |
log_model="all" の場合はすべてのモデルを記録し、log_model=True の場合はトレーニングの最後に記録します |
save_dir |
データが保存されるパス |
ハイパーパラメーターを記録する
追加の設定パラメータを記録する
勾配、パラメータヒストグラム、モデルトポロジーを記録する
モデルのオブジェクトを wandblogger.watch()
に渡すことで、トレーニング中のモデルの勾配とパラメータを監視できます。PyTorch Lightning の WandbLogger
ドキュメントを参照してください。
メトリクスを記録する
WandbLogger
を使用しているときは、LightningModule
内で self.log('my_metric_name', metric_value)
を呼び出すことで W&B にメトリクスを記録できます。たとえば、training_step
や validation_step
メソッド内でこれを行います。
以下のコードスニペットは、メトリクスとハイパーパラメーターを記録するための LightningModule
を定義する方法を示しています。この例では、torchmetrics
ライブラリを使用してメトリクスを計算します。
メトリクスの最小/最大値を記録する
wandb の define_metric
関数を使用して、W&B の要約メトリクスがそのメトリクスの最小、最大、平均、または最良の値を表示するかどうかを定義できます。define_metric
が使用されていない場合、最後に記録された値が要約メトリクスに表示されます。詳細な define_metric
の ドキュメントはこちら と ガイドはこちら を参照してください。
W&B の要約メトリクスで最大の検証精度を追跡するよう W&B に指示するには、トレーニングの開始時に一度だけ wandb.define_metric
を呼び出します:
モデルチェックポイントを作成する
モデルのチェックポイントを W&B の Artifacts として保存するには、Lightning の ModelCheckpoint
コールバックを使用し、WandbLogger
の log_model
引数を設定します。
最新 及び 最良 のエイリアスは、W&B の Artifact からモデルのチェックポイントを簡単に取得できるように自動的に設定されます:
記録されたモデルのチェックポイントは W&B Artifacts UI を通じて表示可能で、完全なモデルリネージも含まれます(UIでのモデルチェックポイントの例はこちら (https://wandb.ai/wandb/arttest/artifacts/model/iv3_trained/5334ab69740f9dda4fed/lineage?_gl=1yyql5q_gaMTQxOTYyNzExOS4xNjg0NDYyNzk1_ga_JH1SJHJQXJ*MTY5MjMwNzI2Mi4yNjkuMS4xNjkyMzA5NjM2LjM3LjAuMA..))。
最良のモデルチェックポイントをブックマークし、チーム全体でそれらを一元化するために、W&B Model Registry にリンクすることができます。
これにより、タスクごとに最良のモデルを整理し、モデルのライフサイクルを管理し、MLライフサイクル全体で簡単な追跡と監査を可能にし、Webhooksやジョブでのダウンストリームアクションを自動化することができます。
画像やテキストなどを記録する
WandbLogger
は、メディアを記録するための log_image
、log_text
、log_table
メソッドを持っています。
他にも、音声、分子、ポイントクラウド、3Dオブジェクトなどのメディアタイプを記録するために、直接的に wandb.log
や trainer.logger.experiment.log
を呼び出すことができます。
Lightning のコールバックシステムを使用して、WandbLogger を介して Weights & Biases にログを記録するタイミングを制御することができます。この例では、検証画像と予測のサンプルをログします:
複数の GPU を使用して Lightning と W&B を使用する
PyTorch Lightning は DDP インターフェースを通じてマルチGPUをサポートしています。ただし、PyTorch Lightning のデザインは GPU をインスタンス化する際に注意が必要です。
Lightning は、トレーニングループ内の各 GPU (またはランク) がまったく同じ方法で、同じ初期条件でインスタンス化されなければならないと仮定しています。ただし、ランク0のプロセスだけが wandb.run
オブジェクトに アクセスでき、非ゼロランクのプロセスには wandb.run = None
となります。これが原因で、非ゼロプロセスが失敗する可能性があります。このような状況になると、ランク0のプロセスが非ゼロランクのプロセスに参加を待つことになり、既にクラッシュしてしまうため、デッドロックに陥る可能性があります。
このため、トレーニングコードのセットアップに注意する必要があります。推奨される方法は、コードを wandb.run
オブジェクトに依存しないようにすることです。
例
Colab のビデオチュートリアルに従うことができます。こちら をクリックしてください。
よくある質問 (FAQ)
W&B は Lightning とどのように統合されていますか?
コアなインテグレーションは、Lightning loggers
API に基づいており、ログのコードをフレームワークに依存しない方法で多く書かせることができます。Logger
は Lightning Trainer
に渡され、この API の豊富な フックとコールバックシステム に基づいてトリガーされます。これにより、研究コードがエンジニアリングやログのコードと完全に分離されます。
追加のコードなしでインテグレーションがログする内容は?
モデルのチェックポイントを W&B に保存し、今後のRunsで使用するために閲覧またはダウンロードできるようにします。また、GPU使用量やネットワークI/Oなどのシステムメトリクス、ハードウェア情報やOS情報などの環境情報、gitコミットやdiffパッチ、ノートブックコンテンツやセッション履歴を含むコードの状態、標準出力に印刷されるものをキャプチャします。
トレーニングセットアップで wandb.run
を使用する必要がある場合はどうすればいいですか?
アクセスが必要な変数のスコープを自分で拡張する必要があります。言い換えれば、初期条件がすべてのプロセスで同じであることを確認してください。
条件が同じならば、os.environ["WANDB_DIR"]
を使用してモデルのチェックポイントディレクトリをセットアップできます。これにより、非ゼロランクプロセスでも wandb.run.dir
にアクセスできます。
フィードバック
このページは役に立ちましたか?
Glad to hear it! If you have further feedback, please let us know.
Sorry to hear that. Please tell us how we can improve.