メインコンテンツまでスキップ

PyTorch Lightning

コラボノートブックで試す →

PyTorch Lightning は、PyTorch コードを整理し、分散トレーニングや16ビット精度などの高度な機能を簡単に追加するための軽量ラッパーを提供します。W&B は、あなたの機械学習実験を記録するための軽量ラッパーを提供しますが、両者を自分で組み合わせる必要はありません。Weights & Biases は PyTorch Lightning ライブラリに直接組み込まれており、WandbLogger 経由で利用できます。

⚡ 数行で素早く始めることができます。​

from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch import Trainer

wandb_logger = WandbLogger(log_model="all")
trainer = Trainer(logger=wandb_logger)
備考

wandb.log() の使用について: WandbLogger は Trainer の global_step を使用して W&B にログを記録します。コード内で直接 wandb.log を追加で呼び出す場合、wandb.log() 内で step 引数を使用しないでください。

代わりに、次のように Trainer の global_step を他のメトリクスと同様にログに記録します：

wandb.log({"accuracy":0.99, "trainer/global_step": step})

どこからでもアクセス可能なインタラクティブダッシュボードとその他の機能

wandb へのサインアップとログイン​

a) 無料アカウントにサインアップします。

b) wandb ライブラリを Pip でインストールします。

c) トレーニングスクリプトでログインするには、www.wandb.ai でアカウントにサインインしている必要があります。その後、Authorizeページ で API キーを見つけることができます。

Weights & Biases を初めて使用する場合、クイックスタート を確認してください。

pip install wandb

wandb login

PyTorch Lightning の WandbLogger の使用​

PyTorch Lightning には複数の WandbLogger (Pytorch) (Fabric) クラスがあり、メトリクス、モデルの重み、メディアなどをシームレスにログできます。単に WandbLogger をインスタンス化し、Lightning の Trainer または Fabric に渡します。

wandb_logger = WandbLogger()
trainer = Trainer(logger=wandb_logger)

Logger の引数​

以下に WandbLogger で最もよく使用されるパラメータを示します。詳細と説明は PyTorch Lightning のドキュメントを参照してください。

パラメータ説明
projectログを記録する wandb プロジェクトを定義
namewandb run に名前を付ける
log_modellog_model="all" の場合は全てのモデルを、log_model=True の場合はトレーニング終了時にモデルをログします
save_dirデータを保存するパス

ハイパーパラメータのログ​

class LitModule(LightningModule):
    def __init__(self, *args, **kwarg):
        self.save_hyperparameters()

追加のコンフィグパラメータをログに記録​

# パラメータを1つ追加
wandb_logger.experiment.config["key"] = value

# 複数のパラメータを追加
wandb_logger.experiment.config.update({key1: val1, key2: val2})

# 直接 wandb モジュールを使用
wandb.config["key"] = value
wandb.config.update()

勾配、パラメータのヒストグラム、モデルのトポロジーをログ​

モデルオブジェクトを wandblogger.watch() に渡すことで、トレーニング中にモデルの勾配とパラメータを監視できます。PyTorch Lightning の WandbLogger ドキュメントを参照してください。

メトリクスのログ​

training_step や validation_step メソッド内で self.log('my_metric_name', metric_vale) と呼び出すことで、LightningModule 内でメトリクスを W&B にログアップすることができます。

以下のコードスニペットは、メトリクスと LightningModule のハイパーパラメータをログする方法を示しています。この例では、メトリクスを計算するために torchmetrics ライブラリを使用します。

import torch
from torch.nn import Linear, CrossEntropyLoss, functional as F
from torch.optim import Adam
from torchmetrics.functional import accuracy
from lightning.pytorch import LightningModule


class My_LitModule(LightningModule):
    def __init__(self, n_classes=10, n_layer_1=128, n_layer_2=256, lr=1e-3):
        """モデルパラメータを定義するためのメソッド"""
        super().__init__()

        # MNIST 画像は (1, 28, 28) (チャンネル, 幅, 高さ)
        self.layer_1 = Linear(28 * 28, n_layer_1)
        self.layer_2 = Linear(n_layer_1, n_layer_2)
        self.layer_3 = Linear(n_layer_2, n_classes)

        self.loss = CrossEntropyLoss()
        self.lr = lr

        # ハイパーパラメータを self.hparams に保存（W&B によって自動的にログされます）
        self.save_hyperparameters()

    def forward(self, x):
        """推論入力 -> 出力に使用されるメソッド"""

        # (b, 1, 28, 28) -> (b, 1*28*28)
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)

        # 3 x (線形 + ReLU) を行う
        x = F.relu(self.layer_1(x))
        x = F.relu(self.layer_2(x))
        x = self.layer_3(x)
        return x

    def training_step(self, batch, batch_idx):
        """1つのバッチからのロスを返す必要があります"""
        _, loss, acc = self._get_preds_loss_accuracy(batch)

        # ロスとメトリクスをログに記録
        self.log("train_loss", loss)
        self.log("train_accuracy", acc)
        return loss

    def validation_step(self, batch, batch_idx):
        """メトリクスをログするために使用されます"""
        preds, loss, acc = self._get_preds_loss_accuracy(batch)

        # ロスとメトリクスをログに記録
        self.log("val_loss", loss)
        self.log("val_accuracy", acc)
        return preds

    def configure_optimizers(self):
        """モデルオプティマイザーを定義する"""
        return Adam(self.parameters(), lr=self.lr)

    def _get_preds_loss_accuracy(self, batch):
        """train/valid/test ステップが似ているための便利な関数"""
        x, y = batch
        logits = self(x)
        preds = torch.argmax(logits, dim=1)
        loss = self.loss(logits, y)
        acc = accuracy(preds, y)
        return preds, loss, acc

メトリクスの最小/最大をログ​

wandb の define_metric 関数を使用して、メトリクスの最小、最大、平均または最良の値を W&B サマリーメトリクスに表示するかどうかを定義できます。define_metric が使用されていない場合、最後にログされた値がサマリーメトリクスに表示されます。define_metric の参考文書はこちら、およびガイドはこちら を参照ください。

W&B サマリーメトリクスで最大検証精度をトラッキングするためには、トレーニング開始時に次のように wandb.define_metric を1回呼ぶだけです：

class My_LitModule(LightningModule):
    ...

    def validation_step(self, batch, batch_idx):
        if trainer.global_step == 0:
            wandb.define_metric("val_accuracy", summary="max")

        preds, loss, acc = self._get_preds_loss_accuracy(batch)

        # ロスとメトリクスをログ
        self.log("val_loss", loss)
        self.log("val_accuracy", acc)
        return preds

モデルチェックポイント​

モデルチェックポイントを W&B Artifacts として保存するには、Lightning の ModelCheckpoint コールバックを使用し、WandbLogger の log_model 引数を設定します：

# `val_accuracy` が増加した場合のみモデルをログ
wandb_logger = WandbLogger(log_model="all")
checkpoint_callback = ModelCheckpoint(monitor="val_accuracy", mode="max")
trainer = Trainer(logger=wandb_logger, callbacks=[checkpoint_callback])

最新と最良のエイリアスが自動的に設定され、W&B Artifact からモデルチェックポイントを簡単に取得できます：

# artifacts パネルで参照を取得できます
# "VERSION" はバージョン（例：「v2」）またはエイリアス（例：「latest」または「best」）のいずれかです
checkpoint_reference = "USER/PROJECT/MODEL-RUN_ID:VERSION"
# チェックポイントをローカルにダウンロード（既にキャッシュされていない場合）
wandb_logger.download_artifact(checkpoint_reference, artifact_type="model")
# チェックポイントをロード
model = LitModule.load_from_checkpoint(Path(artifact_dir) / "model.ckpt")

ログしたモデルチェックポイントは W&B Artifacts UI を通じて表示・ダウンロードでき、完全なモデルリネージを含みます（UI 内のモデルチェックポイントの例 こちら）。

最良のモデルチェックポイントをブックマークしチーム全体で中央管理するために、W&B Model Registry にリンクできます。

ここでは、タスク別に最良のモデルを整理し、モデルのライフサイクルを管理し、MLライフサイクル全体を通じて簡単なトラッキングと監査を実現し、Webhookやジョブを使用して自動化することができます。

画像、テキストなどをログ​

WandbLogger には log_image、log_text、log_table メソッドがあり、メディアをログすることができます。

また、他のメディアタイプ（音声、分子、ポイントクラウド、3Dオブジェクトなど）をログするために、直接 wandb.log または trainer.logger.experiment.log を呼び出すこともできます。

# テンソル、numpy配列、またはPIL画像を使用
wandb_logger.log_image(key="samples", images=[img1, img2])

# キャプションを追加
wandb_logger.log_image(key="samples", images=[img1, img2], caption=["tree", "person"])

# ファイルパスを使用
wandb_logger.log_image(key="samples", images=["img_1.jpg", "img_2.jpg"])

# トレーナーで .log を使用
trainer.logger.experiment.log(
    {"samples": [wandb.Image(img, caption=caption) for (img, caption) in my_images]},
    step=current_trainer_global_step,
)

Lightning の Callbacks システムを使用して、WandbLogger を介して Weights & Biases にログを記録するタイミングを制御できます。この例では、検証画像と予測のサンプルをログに記録します：

import torch
import wandb
import lightning.pytorch as pl
from lightning.pytorch.loggers import WandbLogger

# または
# from wandb.integration.lightning.fabric import WandbLogger

class LogPredictionSamplesCallback(Callback):
    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx
    ):
        """検証バッチが終了したときに呼び出されます。"""

        # `outputs` は `LightningModule.validation_step` からのもので、
        # この場合は私たちのモデルの予測に相当します。

        # 最初のバッチから20サンプル画像の予測をログ
        if batch_idx == 0:
            n = 20
            x, y = batch
            images = [img for img in x[:n]]
            captions = [
                f"Ground Truth: {y_i} - Prediction: {y_pred}"
                for y_i, y_pred in zip(y[:n], outputs[:n])
            ]

            # オプション1: `WandbLogger.log_image` で画像をログ
            wandb_logger.log_image(key="sample_images", images=images, caption=captions)

            # オプション2: 画像と予測をW&Bテーブルとしてログ
            columns = ["image", "ground truth", "prediction"]
            data = [
                [wandb.Image(x_i), y_i, y_pred] or x_i,
                y_i,
                y_pred in list(zip(x[:n], y[:n], outputs[:n])),
            ]
            wandb_logger.log_table(key="sample_table", columns=columns, data=data)


trainer = pl.Trainer(callbacks=[LogPredictionSamplesCallback()])

LightningとW&Bを使用して複数のGPUをどのように使うか？​

PyTorch Lightning には DDP インターフェースを通じたマルチGPUサポートがあります。ただし、PyTorch Lightning の設計には、GPU（またはランク）のトレーニングループが全く同じ初期条件でインスタンス化されることを前提としています。しかし、ランク0プロセスだけが wandb.run オブジェクトにアクセスでき、ランク0以外のプロセスは wandb.run = None になります。このような状況は、ランク0プロセスが他のランクのプロセスを待ち続けるため、デッドロックに陥る可能性があります。

したがって、トレーニングコードのセットアップに注意が必要です。推奨される方法は、コードを wandb.run オブジェクトに依存しないように設定することです。

class MNISTClassifier(pl.LightningModule):
    def __init__(self):
        super(MNISTClassifier, self).__init__()

        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

        self.loss = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = self.loss(y_hat, y)

        self.log("train/loss", loss)
        return {"train_loss": loss}

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = self.loss(y_hat, y)

        self.log("val/loss", loss)
        return {"val_loss": loss}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)


def main():
    # すべてのランダムシードを同じ値にする設定
    # これは分散トレーニング環境で重要です
    # 各ランクは独自の初期重みを取得します。
    # それらが一致しないと、勾配も一致せずに
    # トレーニングが収束しない可能性があります。
    pl.seed_everything(1)

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4)

    model = MNISTClassifier()
    wandb_logger = WandbLogger(project="<project_name>")
    callbacks = [
        ModelCheckpoint(
            dirpath="checkpoints",
            every_n_train_steps=100,
        ),
    ]
    trainer = pl.Trainer(
        max_epochs=3, gpus=2, logger=wandb_logger, strategy="ddp", callbacks=callbacks
    )
    trainer.fit(model, train_loader, val_loader)

インタラクティブな例を確認！​

ビデオチュートリアルに沿って進めるか、こちら のチュートリアルコラボで進めることができます。

よくある質問​

W&B はどのように Lightning に統合されますか？​

コア統合はLightning loggers API に基づいており、多くのログ記録コードをフレームワークに依存しない方法で記述することができます。Logger は Lightning Trainer に渡され、その API の豊富なフックおよびコールバックシステム に基づいてトリガーされます。これにより、研究コードがエンジニアリングおよびログ記録コードから分離されます。

追加のコードなく統合は何をログしますか？​

モデルのチェックポイントをW&Bに保存し、将来のランで使用するために表示またはダウンロードできます。また、システムメトリクス（GPU使用率やネットワークI/Oなど）、ハードウェアやOS情報などの環境情報、コード状態（gitコミット情報やdiffパッチ、ノートブックの内容やセッション履歴を含む）、標準出力に印刷される全てをキャプチャします。

トレーニングセットアップで wandb.run をどうしても使用する必要がある場合はどうすればよいですか？​

必要な変数のスコープを自分で拡大することになります。つまり、全てのプロセスで同じ初期条件を設定することを確認してください。

if os.environ.get("LOCAL_RANK", None) is None:
    os.environ["WANDB_DIR"] = wandb.run.dir
On this page
Was this page helpful?👍👎