> ## Documentation Index > Fetch the complete documentation index at: https://wb-21fd5541-serverless-sft-revamp.mintlify.site/llms.txt > Use this file to discover all available pages before exploring further. # PyTorch Lightning export const ColabLink = ({url}) => Try in Colab ; PyTorch Lightning を使用して画像分類パイプラインを構築します。コードの可読性と再現性を高めるために、この [スタイルガイド](https://lightning.ai/docs/pytorch/stable/starter/style_guide.html) に従います。これに関する詳しい解説は [こちら](https://wandb.ai/wandb/wandb-lightning/reports/Image-Classification-using-PyTorch-Lightning--VmlldzoyODk1NzY) で確認できます。 ## PyTorch Lightning と W\&B のセットアップこのチュートリアルでは、PyTorch Lightning と W\&B が必要です。 ```shell theme={null} pip install lightning -q pip install wandb -qU ``` ```python theme={null} import lightning.pytorch as pl # お気に入りの機械学習トラッキングツール from lightning.pytorch.loggers import WandbLogger import torch from torch import nn from torch.nn import functional as F from torch.utils.data import random_split, DataLoader from torchmetrics import Accuracy from torchvision import transforms from torchvision.datasets import CIFAR10 import wandb ``` 次に、wandb アカウントにログインする必要があります。 ``` wandb.login() ``` ## DataModule - 私たちが求めるデータパイプライン DataModule は、データ関連のフックを LightningModule から切り離す方法であり、データセットに依存しないモデルを開発できるようにします。これにより、データパイプラインを共有可能で再利用可能な一つのクラスに整理できます。DataModule は、PyTorch におけるデータプロセッシングの 5 つのステップをカプセル化します： * ダウンロード / トークン化 / 処理 * クリーニング、および（必要に応じて）ディスクへの保存 * Dataset へのロード * 変換（回転、トークン化など）の適用 * DataLoader へのラップ DataModule についての詳細は [こちら](https://lightning.ai/docs/pytorch/stable/data/datamodule.html) をご覧ください。CIFAR-10 データセット用の DataModule を構築してみましょう。 ```python theme={null} class CIFAR10DataModule(pl.LightningDataModule): def __init__(self, batch_size, data_dir: str = './'): super().__init__() self.data_dir = data_dir self.batch_size = batch_size self.transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) self.num_classes = 10 def prepare_data(self): CIFAR10(self.data_dir, train=True, download=True) CIFAR10(self.data_dir, train=False, download=True) def setup(self, stage=None): # データローダーで使用するトレーニング/検証用データセットを割り当て if stage == 'fit' or stage is None: cifar_full = CIFAR10(self.data_dir, train=True, transform=self.transform) self.cifar_train, self.cifar_val = random_split(cifar_full, [45000, 5000]) # データローダーで使用するテスト用データセットを割り当て if stage == 'test' or stage is None: self.cifar_test = CIFAR10(self.data_dir, train=False, transform=self.transform) def train_dataloader(self): return DataLoader(self.cifar_train, batch_size=self.batch_size, shuffle=True) def val_dataloader(self): return DataLoader(self.cifar_val, batch_size=self.batch_size) def test_dataloader(self): return DataLoader(self.cifar_test, batch_size=self.batch_size) ``` ## Callbacks コールバックは、プロジェクト間で再利用可能な自己完結型のプログラムです。PyTorch Lightning には、よく使用されるいくつかの [組み込みコールバック](https://lightning.ai/docs/pytorch/latest/extensions/callbacks.html#built-in-callbacks) が用意されています。 PyTorch Lightning のコールバックについての詳細は [こちら](https://lightning.ai/docs/pytorch/latest/extensions/callbacks.html) をご覧ください。 ### 組み込みコールバックこのチュートリアルでは、[Early Stopping](https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.EarlyStopping.html#lightning.callbacks.EarlyStopping) と [Model Checkpoint](https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.ModelCheckpoint.html#pytorch_lightning.callbacks.ModelCheckpoint) の組み込みコールバックを使用します。これらは `Trainer` に渡すことができます。 ### カスタムコールバック Keras のカスタムコールバックに慣れているなら、PyTorch のパイプラインでも同じことができるのは非常に魅力的です。今回は画像分類を行っているので、いくつかの画像サンプルに対するモデルの予測を可視化できると便利です。これをコールバックの形式にすることで、早い段階でのモデルのデバッグに役立ちます。 ```python theme={null} class ImagePredictionLogger(pl.callbacks.Callback): def __init__(self, val_samples, num_samples=32): super().__init__() self.num_samples = num_samples self.val_imgs, self.val_labels = val_samples def on_validation_epoch_end(self, trainer, pl_module): # テンソルを CPU に移動 val_imgs = self.val_imgs.to(device=pl_module.device) val_labels = self.val_labels.to(device=pl_module.device) # モデルの予測を取得 logits = pl_module(val_imgs) preds = torch.argmax(logits, -1) # 画像を wandb Image としてログ trainer.logger.experiment.log({ "examples":[wandb.Image(x, caption=f"Pred:{pred}, Label:{y}") for x, pred, y in zip(val_imgs[:self.num_samples], preds[:self.num_samples], val_labels[:self.num_samples])] }) ``` ## LightningModule - システムの定義 LightningModule はモデルではなく「システム」を定義します。ここでは、すべての研究コードを一つのクラスにグループ化し、自己完結型にすることをシステムと呼びます。`LightningModule` は、PyTorch のコードを以下の 5 つのセクションに整理します： * 計算（`__init__`） * トレーニングループ（`training_step`） * 検証ループ（`validation_step`） * テストループ（`test_step`） * オプティマイザー（`configure_optimizers`）これにより、簡単に共有可能なデータセットに依存しないモデルを構築できます。CIFAR-10 分類用のシステムを構築しましょう。 ```python theme={null} class LitModel(pl.LightningModule): def __init__(self, input_shape, num_classes, learning_rate=2e-4): super().__init__() # ハイパーパラメーターのログ記録 self.save_hyperparameters() self.learning_rate = learning_rate self.conv1 = nn.Conv2d(3, 32, 3, 1) self.conv2 = nn.Conv2d(32, 32, 3, 1) self.conv3 = nn.Conv2d(32, 64, 3, 1) self.conv4 = nn.Conv2d(64, 64, 3, 1) self.pool1 = torch.nn.MaxPool2d(2) self.pool2 = torch.nn.MaxPool2d(2) n_sizes = self._get_conv_output(input_shape) self.fc1 = nn.Linear(n_sizes, 512) self.fc2 = nn.Linear(512, 128) self.fc3 = nn.Linear(128, num_classes) self.accuracy = Accuracy(task='multiclass', num_classes=num_classes) # conv ブロックから Linear レイヤーに入る出力テンソルのサイズを返す def _get_conv_output(self, shape): batch_size = 1 input = torch.autograd.Variable(torch.rand(batch_size, *shape)) output_feat = self._forward_features(input) n_size = output_feat.data.view(batch_size, -1).size(1) return n_size # conv ブロックからの特徴テンソルを返す def _forward_features(self, x): x = F.relu(self.conv1(x)) x = self.pool1(F.relu(self.conv2(x))) x = F.relu(self.conv3(x)) x = self.pool2(F.relu(self.conv4(x))) return x # 推論時に使用 def forward(self, x): x = self._forward_features(x) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = F.log_softmax(self.fc3(x), dim=1) return x def training_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = F.nll_loss(logits, y) # トレーニングメトリクス preds = torch.argmax(logits, dim=1) acc = self.accuracy(preds, y) self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True) self.log('train_acc', acc, on_step=True, on_epoch=True, logger=True) return loss def validation_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = F.nll_loss(logits, y) # 検証メトリクス preds = torch.argmax(logits, dim=1) acc = self.accuracy(preds, y) self.log('val_loss', loss, prog_bar=True) self.log('val_acc', acc, prog_bar=True) return loss def test_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = F.nll_loss(logits, y) # 検証メトリクス preds = torch.argmax(logits, dim=1) acc = self.accuracy(preds, y) self.log('test_loss', loss, prog_bar=True) self.log('test_acc', acc, prog_bar=True) return loss def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate) return optimizer ``` ## トレーニングと評価 `DataModule` を使ってデータパイプラインを、`LightningModule` を使ってモデルアーキテクチャーとトレーニングループを整理しました。あとは PyTorch Lightning の `Trainer` がすべてを自動化してくれます。 Trainer が自動化するもの： * エポックとバッチのイテレーション * `optimizer.step()`、`backward`、`zero_grad()` の呼び出し * `.eval()` の呼び出し、勾配の有効化/無効化 * 重みの保存とロード * W\&B へのログ記録 * マルチ GPU トレーニングのサポート * TPU のサポート * 16-bit トレーニングのサポート ```python theme={null} dm = CIFAR10DataModule(batch_size=32) # データローダーにアクセスするために prepare_data と setup を呼び出す必要があります。 dm.prepare_data() dm.setup() # カスタム ImagePredictionLogger コールバックが画像予測をログするために必要なサンプル val_samples = next(iter(dm.val_dataloader())) val_imgs, val_labels = val_samples[0], val_samples[1] val_imgs.shape, val_labels.shape ``` ```python theme={null} model = LitModel((3, 32, 32), dm.num_classes) # wandb logger の初期化 wandb_logger = WandbLogger(project='wandb-lightning', job_type='train') # コールバックの初期化 early_stop_callback = pl.callbacks.EarlyStopping(monitor="val_loss") checkpoint_callback = pl.callbacks.ModelCheckpoint() # trainer の初期化 trainer = pl.Trainer(max_epochs=2, logger=wandb_logger, callbacks=[early_stop_callback, ImagePredictionLogger(val_samples), checkpoint_callback], ) # モデルのトレーニング trainer.fit(model, dm) # ホールドアウトされたテストセットでモデルを評価 ⚡⚡ trainer.test(dataloaders=dm.test_dataloader()) # wandb run を終了 run.finish() ``` ## 最後に私は TensorFlow/Keras エコシステムの出身で、PyTorch は洗練されたフレームワークであるものの、少し敷居が高いと感じていました（あくまで個人的な経験です）。しかし、PyTorch Lightning を試してみると、PyTorch から遠ざかっていた理由のほとんどが解消されていることに気づきました。私の感動を簡単にまとめます： * 以前：従来の PyTorch のモデル定義はバラバラになりがちでした。モデルは `model.py` スクリプトにあり、トレーニングループは `train.py` ファイルにあるといった具合です。パイプラインを理解するために何度もファイルを行き来する必要がありました。 * 現在：`LightningModule` がシステムとして機能し、モデルが `training_step` や `validation_step` などと一緒に定義されます。モジュール化され、共有しやすくなりました。 * 以前：TensorFlow/Keras の最大の魅力は入力データパイプラインでした。データセットカタログが豊富で成長し続けています。PyTorch のデータパイプラインは最大の懸念点でした。通常の PyTorch コードでは、データのダウンロード/クリーニング/準備が多くのファイルに分散していることがよくあります。 * 現在：DataModule がデータパイプラインを一つの共有・再利用可能なクラスにまとめます。これは単に `train_dataloader`、`val_dataloader`、`test_dataloader` と、それに対応する変換やデータ処理/ダウンロードステップの集まりにすぎません。 * 以前：Keras では、`model.fit` でトレーニング、`model.predict` で推論を実行できました。`model.evaluate` はテストデータに対するシンプルで使い勝手の良い評価を提供していました。PyTorch ではそうはいかず、通常は別々の `train.py` と `test.py` ファイルが必要でした。 * 現在：`LightningModule` を導入すれば、`Trainer` がすべてを自動化します。`trainer.fit` と `trainer.test` を呼び出すだけで、モデルのトレーニングと評価が完了します。 * 以前：TensorFlow は TPU が大好きですが、PyTorch は... * 現在：PyTorch Lightning を使えば、同じモデルを複数の GPU や TPU でトレーニングするのが非常に簡単です。 * 以前：私はコールバックの大ファンで、カスタムコールバックを書くのが好きです。Early Stopping のような些細なことでさえ、従来の PyTorch では議論の的になることがありました。 * 現在：PyTorch Lightning では Early Stopping や Model Checkpointing を使うのは非常に簡単です。カスタムコールバックを書くこともできます。 ## 🎨 結論とリソースこのレポートがお役に立てば幸いです。ぜひコードを動かして、お好みのデータセットで画像分類器をトレーニングしてみてください。 PyTorch Lightning についてさらに詳しく学ぶためのリソースをいくつか紹介します： * [ステップバイステップガイド](https://lightning.ai/docs/pytorch/latest/starter/introduction.html): 公式チュートリアルの一つです。ドキュメントが非常によく書かれており、学習リソースとして強くお勧めします。 * [Pytorch Lightning を W\&B で使用する](https://wandb.me/lightning): PyTorch Lightning で W\&B を使用する方法を学べるクイックな Colab です。