PyTorch Lightning

組み込みの WandbLogger を使用して、PyTorch Lightning で W&B による実験管理とモデル チェックポイントを行います。

PyTorch Lightning は、PyTorch コードを整理し、分散トレーニングや 16 ビット精度などの高度な機能を簡単に追加できる軽量なラッパーを提供します。W&B は、ML 実験をログするための軽量なラッパーを提供します。しかし、この 2 つを自分で組み合わせる必要はありません。W&B は、WandbLogger を通じて PyTorch Lightning ライブラリに直接統合されています。

​

Lightningと統合する

PyTorch Logger

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})

Fabric Logger

import lightning as L
from wandb.integration.lightning.fabric import WandbLogger

wandb_logger = WandbLogger(log_model="all")
fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()
fabric.log_dict({"important_metric": important_metric})

​

サインアップしてAPIキーを発行する

APIキーを使うと、お使いのマシンをW&Bに対して認証できます。APIキーはプロフィールから発行できます。

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

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

​

wandb ライブラリをインストールしてログインする

wandb ライブラリをローカル環境にインストールしてログインするには、次の手順を実行します。

コマンドライン

1. WANDB_API_KEY 環境変数に APIキー を設定します。

   export WANDB_API_KEY=<your_api_key>
   

2. wandb ライブラリをインストールし、ログインします。

   pip install wandb
   
   wandb login
   

Python

pip install wandb

import wandb
wandb.login()

Python notebook

!pip install wandb

import wandb
wandb.login()

​

PyTorch Lightning の WandbLogger を使用する

PyTorch Lightning には、メトリクス、モデルの重み、メディアなどをログするための WandbLogger クラスが複数あります。

• PyTorch
• Fabric

Lightning と統合するには、WandbLogger をインスタンス化して、Lightning の Trainer または Fabric に渡します。

PyTorch Logger

trainer = Trainer(logger=wandb_logger)

Fabric Logger

fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()
fabric.log_dict({
    "important_metric": important_metric
})

​

よく使われるロガー引数

以下は、WandbLogger でよく使われるパラメーターの一部です。ロガー引数の詳細については、PyTorch Lightning のドキュメントを参照してください。

• PyTorch
• Fabric

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

​

ハイパーパラメーターをログする

PyTorch Logger

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

Fabric Logger

wandb_logger.log_hyperparams(
    {
        "hyperparameter_1": hyperparameter_1,
        "hyperparameter_2": hyperparameter_2,
    }
)

​

追加の設定パラメーターをログする

# パラメーターを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 ドキュメントを参照してください

​

メトリクスをログする

PyTorch Logger

WandbLogger を使用している場合は、LightningModule の training_step や validation_step メソッド内で self.log('my_metric_name', metric_vale) を呼び出すことで、メトリクスを W&B にログできます。

以下のコードスニペットは、メトリクスと LightningModule のハイパーパラメーターをログするように 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 images は (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):
        """推論の input -> output に使用するメソッド"""

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

        # 3 x (linear + 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):
        """単一バッチから損失を返す必要があります"""
        _, 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

Fabric Logger

import lightning as L
import torch
import torchvision as tv
from wandb.integration.lightning.fabric import WandbLogger
import wandb

fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()

model = tv.models.resnet18()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
model, optimizer = fabric.setup(model, optimizer)

train_dataloader = fabric.setup_dataloaders(
torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
)

model.train()
for epoch in range(num_epochs):
for batch in train_dataloader:
optimizer.zero_grad()
loss = model(batch)
loss.backward()
optimizer.step()
fabric.log_dict({"loss": loss})

​

メトリクスの最小値/最大値をログする

wandb の define_metric 関数を使用すると、W&B summary メトリクスにそのメトリクスの最小値、最大値、平均値、または最良値のどれを表示するかを定義できます。define_metric _ を使用しない場合は、最後にログされた値が summary メトリクスに表示されます。詳細については、define_metric のリファレンスドキュメントとガイドを参照してください。

W&B summary メトリクスで検証 accuracy の最大値をトラッキングするには、トレーニングの開始時に wandb.define_metric() を一度だけ call します。

PyTorch Logger

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

Fabric Logger

wandb.define_metric("val_accuracy", summary="max")
fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()
fabric.log_dict({"val_accuracy": val_accuracy})

​

モデルのチェックポイントを保存する

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

PyTorch Loggerpython trainer = Trainer(logger=wandb_logger, callbacks=[checkpoint_callback])

Fabric Loggerpython fabric = L.Fabric(loggers=[wandb_logger], callbacks=[checkpoint_callback])

latest と best のエイリアスは自動的に設定されるため、W&B Artifact からモデル チェックポイントを簡単に取得できます。

# 参照は Artifacts パネルで取得できます
# "VERSION" には、バージョン（例: "v2"）またはエイリアス（"latest" または "best"）を指定できます
checkpoint_reference = "USER/PROJECT/MODEL-RUN_ID:VERSION"

Logger 経由python # チェックポイントをローカルにダウンロードします（未キャッシュの場合） wandb_logger.download_artifact(checkpoint_reference, artifact_type="model")

wandb 経由

# チェックポイントをローカルにダウンロードします（未キャッシュの場合） run = wandb.init(project="MNIST") artifact = run.use_artifact(checkpoint_reference, type="model") artifact_dir = artifact.download()

PyTorch Logger

# チェックポイントを読み込みます
model = LitModule.load_from_checkpoint(Path(artifact_dir) / "model.ckpt")

Fabric Logger

# 生のチェックポイントを取得します
full_checkpoint = fabric.load(Path(artifact_dir) / "model.ckpt")

model.load_state_dict(full_checkpoint["model"])
optimizer.load_state_dict(full_checkpoint["optimizer"])

ログしたモデル チェックポイントは W&B Artifacts UI から確認でき、モデルの完全なリネージ情報も含まれます (UI でのモデル チェックポイントの例は こちら を参照してください) 。

ベストなモデル チェックポイントをブックマークしてチーム全体で一元管理するには、W&B Model Registry にリンクできます。

ここでは、ベストなモデルをタスク別に整理し、モデルのライフサイクルを管理し、ML ライフサイクル全体を通じて容易にトラッキングや監査を行い、webhook やジョブで下流のアクションを自動化できます。

​

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

WandbLogger には、メディアをログするための log_image、log_text、log_table メソッドがあります。

wandb.log() または trainer.logger.experiment.log() を直接呼び出して、オーディオ、分子、ポイントクラウド、3D オブジェクトなど、ほかのメディアタイプをログすることもできます。

画像をログする

# tensors、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"])

# trainer で .log を使用

trainer.logger.experiment.log(
{"samples": [wandb.Image(img, caption=caption) for (img, caption) in my_images]},
step=current_trainer_global_step,
)

テキストをログする

# data はリストのリストである必要があります
columns = ["input", "label", "prediction"]
my_data = [["cheese", "english", "english"], ["fromage", "french", "spanish"]]

# columns と data を使用
wandb_logger.log_text(key="my_samples", columns=columns, data=my_data)

# pandas DataFrame を使用
wandb_logger.log_text(key="my_samples", dataframe=my_dataframe)

Tables をログする

# テキストのキャプション、画像、オーディオを含む W&B Table をログします
columns = ["caption", "image", "sound"]

# data はリストのリストである必要があります

my_data = [
["cheese", wandb.Image(img_1), wandb.Audio(snd_1)],
["wine", wandb.Image(img_2), wandb.Audio(snd_2)],
]

# Table をログします

wandb_logger.log_table(key="my_samples", columns=columns, data=data)

Lightning のコールバックシステムを使用すると、WandbLogger を介して W&B にいつログするかを制御できます。この例では、検証画像と予測のサンプルをログします。

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 Table としてログする
            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 をどのようにインスタンス化するかに注意が必要です。

Lightning では、トレーニングループ内の各 GPU (または ランク) を、同じ初期条件でまったく同じようにインスタンス化することが前提となっています。一方で、wandb.run オブジェクトにアクセスできるのは ランク 0 のプロセスだけで、0 以外の ランク のプロセスでは wandb.run = None になります。これにより、0 以外のプロセスが失敗する可能性があります。このような状況では、ランク 0 のプロセスは、すでにクラッシュした 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)

​

例

Colab ノートブック付きの動画チュートリアルに沿って進めることができます。

​

よくある質問

​

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

主要なインテグレーションは Lightning loggers API に基づいており、そのおかげでロギングコードの大部分をフレームワーク非依存の形で記述できます。Logger は Lightning Trainer に渡され、この API の充実した hook-and-callback system によってトリガーされます。これにより、研究用コードをエンジニアリングやロギングのコードから明確に分離できます。

​

追加のコードを書かなくても、インテグレーションは何をログしますか？

モデル チェックポイントは W&B に保存され、そこで確認したり、今後の run で使用するためにダウンロードしたりできます。さらに、GPU 使用率やネットワーク I/O などのシステム メトリクス、ハードウェアや OS などの環境情報、コードの状態 (git commit や diff patch、ノートブックの内容、セッション履歴を含む) 、および標準出力に出力された内容も取得します。

​

トレーニング設定で wandb.run を使用する必要がある場合はどうすればよいですか？

アクセスする必要がある変数のスコープは、自分で広げる必要があります。つまり、すべてのプロセスで初期条件が同じになるようにしてください。

if os.environ.get("LOCAL_RANK", None) is None:
    os.environ["WANDB_DIR"] = wandb.run.dir

その場合は、os.environ["WANDB_DIR"] を使用してモデル チェックポイントのディレクトリを設定できます。これにより、ランクが 0 以外の任意のプロセスから wandb.run.dir にアクセスできます。