pytorch-lightning

PyTorch Lightning

概要

PyTorch Lightning は、柔軟性を完全に維持しながら定型コードをなくせるよう、PyTorch コードを整理するディープラーニングフレームワークです。学習ワークフローや複数デバイスのオーケストレーションを自動化し、ニューラルネットワークの学習や複数の GPU／TPU にわたるスケーリングのベストプラクティスを実装できます。

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

次のような場合に、このスキルを使用してください。

• PyTorch Lightning を使用してニューラルネットワークを構築、学習、デプロイする場合
• PyTorch コードを LightningModule に整理する場合
• 複数の GPU／TPU での学習向けに Trainer を設定する場合
• LightningDataModule を使用してデータパイプラインを実装する場合
• コールバック、ロギング、分散学習戦略（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. コールバック - 拡張可能な学習ロジック

LightningModule を変更せずに、特定の学習フックでカスタム機能を追加できます。組み込みコールバックには次のものがあります。

• ModelCheckpoint - 最良／最新のモデルを保存する
• EarlyStopping - 指標が頭打ちになったときに停止する
• LearningRateMonitor - 学習率スケジューラーの変化を追跡する
• BatchSizeFinder - 最適なバッチサイズを自動的に決定する

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

5. ロギング - 実験の追跡

複数のロギングプラットフォームと統合できます。

• 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、小規模なトランスフォーマー向け）
• FSDP - 5億以上のパラメーターを持つモデル（大規模なトランスフォーマー向け。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. データを準備：

   # Option 1: Direct DataLoaders
   train_loader = DataLoader(train_dataset, batch_size=32)
   
   # Option 2: LightningDataModule (recommended for reusability)
   dm = MyDataModule(batch_size=32)
   

3. トレーニング：

   trainer = L.Trainer(max_epochs=10, accelerator="gpu", devices=2)
   trainer.fit(model, train_loader)  # or 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 - 一般的なパターン、ヒント、注意点