PyTorch Lightning

概要

PyTorch Lightning は、ボイラープレートコードを削減しながら完全な柔軟性を保つ、PyTorch コードを整理するディープラーニングフレームワークです。訓練ワークフロー、マルチデバイスオーケストレーション、そして複数の GPU/TPU にわたるニューラルネットワークの訓練とスケーリングのベストプラクティスを自動化します。

このスキルを使用する場合

以下の場合、このスキルを使用してください:

• PyTorch Lightning を使用したニューラルネットワークの構築、訓練、デプロイ
• PyTorch コードを LightningModule に整理する
• マルチ GPU/TPU 訓練のための Trainer を構成する
• LightningDataModules を使用したデータパイプラインの実装
• コールバック、ロギング、分散訓練戦略(DDP、FSDP、DeepSpeed)の使用
• ディープラーニングプロジェクトを専門的に構造化する

コア機能

1. LightningModule - モデル定義

PyTorch モデルを 6 つの論理的なセクションに整理します:

1. 初期化 - __init__() と setup()
2. 訓練ループ - training_step(batch, batch_idx)
3. 検証ループ - validation_step(batch, batch_idx)
4. テストループ - test_step(batch, batch_idx)
5. 予測 - predict_step(batch, batch_idx)
6. オプティマイザー設定 - configure_optimizers()

クイックテンプレートリファレンス: 完全なボイラープレートについては scripts/template_lightning_module.py を参照してください。

詳細なドキュメント: 包括的なメソッドドキュメント、フック、プロパティ、ベストプラクティスについては references/lightning_module.md を参照してください。

2. Trainer - 訓練自動化

Trainer は訓練ループ、デバイス管理、勾配操作、コールバックを自動化します。主な機能:

• DDP、FSDP、DeepSpeed などの戦略選択によるマルチ GPU/TPU サポート
• 自動混合精度訓練
• 勾配の蓄積とクリッピング
• チェックポイント保存と早期停止
• プログレスバーとロギング

クイックセットアップリファレンス: 一般的な Trainer 設定については scripts/quick_trainer_setup.py を参照してください。

詳細なドキュメント: すべてのパラメータ、メソッド、設定オプションについては references/trainer.md を参照してください。

3. LightningDataModule - データパイプライン整理

再利用可能なクラスにすべてのデータ処理ステップをカプセル化します:

1. prepare_data() - データのダウンロードと処理(単一プロセス)
2. setup() - データセットの作成と変換の適用(GPU ごと)
3. train_dataloader() - 訓練 DataLoader を返す
4. val_dataloader() - 検証 DataLoader を返す
5. test_dataloader() - テスト DataLoader を返す

クイックテンプレートリファレンス: 完全なボイラープレートについては scripts/template_datamodule.py を参照してください。

詳細なドキュメント: メソッドの詳細と使用パターンについては references/data_module.md を参照してください。

4. Callbacks - 拡張可能な訓練ロジック

LightningModule を修正することなく、特定の訓練フックでカスタム機能を追加します。組み込みコールバック:

• ModelCheckpoint - 最高/最新モデルを保存
• EarlyStopping - メトリクスがプラトーに達したときに停止
• LearningRateMonitor - LR スケジューラの変化を追跡
• BatchSizeFinder - 最適バッチサイズを自動決定

詳細なドキュメント: 組み込みコールバックとカスタムコールバック作成については references/callbacks.md を参照してください。

5. Logging - 実験追跡

複数のロギングプラットフォームと統合:

• TensorBoard(デフォルト)
• Weights & Biases(WandbLogger)
• MLflow(MLFlowLogger)
• Neptune(NeptuneLogger)
• Comet(CometLogger)
• CSV(CSVLogger)

任意の LightningModule メソッドで self.log("metric_name", value) を使用してメトリクスをログ出力します。

詳細なドキュメント: ロガーのセットアップと設定については references/logging.md を参照してください。

6. 分散訓練 - マルチデバイスへのスケーリング

モデルサイズに基づいて適切な戦略を選択します:

• DDP - 5 億パラメータ未満のモデル用(ResNet、小型 Transformer)
• FSDP - 5 億パラメータ以上のモデル用(大型 Transformer、Lightning ユーザー向け推奨)
• DeepSpeed - 最新機能と細粒度制御用

次のように設定します: Trainer(strategy="ddp", accelerator="gpu", devices=4)

詳細なドキュメント: 戦略の比較と設定については references/distributed_training.md を参照してください。

7. ベストプラクティス

• デバイス非依存のコード - .cuda() の代わりに self.device を使用
• ハイパーパラメータ保存 - __init__() で self.save_hyperparameters() を使用
• メトリクスロギング - デバイス全体での自動集約のために self.log() を使用
• 再現性 - seed_everything() と Trainer(deterministic=True) を使用
• デバッグ - Trainer(fast_dev_run=True) を使用して 1 バッチでテスト

詳細なドキュメント: 一般的なパターンと落とし穴については references/best_practices.md を参照してください。

クイックワークフロー

1. モデルを定義:

   class MyModel(L.LightningModule):
       def __init__(self):
           super().__init__()
           self.save_hyperparameters()
           self.model = YourNetwork()
   
       def training_step(self, batch, batch_idx):
           x, y = batch
           loss = F.cross_entropy(self.model(x), y)
           self.log("train_loss", loss)
           return loss
   
       def configure_optimizers(self):
           return torch.optim.Adam(self.parameters())
   

2. データを準備:

   # オプション 1: 直接 DataLoader
   train_loader = DataLoader(train_dataset, batch_size=32)
   
   # オプション 2: LightningDataModule(再利用性のため推奨)
   dm = MyDataModule(batch_size=32)
   

3. 訓練:

   trainer = L.Trainer(max_epochs=10, accelerator="gpu", devices=2)
   trainer.fit(model, train_loader)  # または trainer.fit(model, datamodule=dm)
   

リソース

scripts/

PyTorch Lightning の一般的なパターン向けの実行可能な Python テンプレート:

• template_lightning_module.py - 完全な LightningModule ボイラープレート
• template_datamodule.py - 完全な LightningDataModule ボイラープレート
• quick_trainer_setup.py - 一般的な Trainer 設定例

references/

各 PyTorch Lightning コンポーネントの詳細なドキュメント:

• lightning_module.md - 包括的な LightningModule ガイド(メソッド、フック、プロパティ)
• trainer.md - Trainer 設定とパラメータ
• data_module.md - LightningDataModule パターンとメソッド
• callbacks.md - 組み込みコールバックとカスタムコールバック
• logging.md - ロガー統合と使用方法
• distributed_training.md - DDP、FSDP、DeepSpeed の比較とセットアップ
• best_practices.md - 一般的なパターン、ヒント、落とし穴