PyTorch Lightning 入門

PyTorch Lightning によるスケーラブルな AI モデルの構築

Sergiy Tkachuk

Director, GenAI Productivity

PyTorch と PyTorch Lightning

標準的な PyTorch:

多くの手作業が必要
明示的なトレーニングループの記述
GPU/TPU の処理、ロギング、チェックポイント管理

PyTorch と PyTorch Lightning

PyTorch Lightning:

PyTorch をベースに構築
自動化対象:
- トレーニング
- チェックポイント管理
- ロギング
ボイラープレートコードを削減
スケーラビリティと再現性を向上

PyTorch Lightning のロゴ

PyTorch Lightning の概要

例: グローバル e コマース企業によるワークフローの効率化
- ビジュアル検索モデルの開発
- トレーニングループの自動化
- 最小限のボイラープレートコードで高速に反復

コアコンポーネント: LightningModule と Trainer

from lightning.pytorch import LightningModule
from lightning.pytorch import Trainer

Lightning の構造

主要コンポーネント:

LightningModule: モデルのコアロジック

ロジックを表す電球

Lightning の構造

主要コンポーネント:

LightningModule: モデルのコアロジック
Lightning Trainer: トレーニングを統括

ロジックを表す電球と処理を表す歯車

Lightning の構造

主要コンポーネント:

LightningModule: モデルのコアロジック
Lightning Trainer: トレーニングを統括
DataModule: データパイプラインを整理
Callbacks: イベントを自動化
Logger: 実験を記録

ロジックを表す電球と処理を表す歯車

LightningModule の活用例

重要なポイント:

__init__: モデルのアーキテクチャを定義
forward(): モデルにデータを通過させる
training_step(): トレーニングを定義
カスタムフックを利用可能

import lightning.pytorch as pl

class LightClassifier(pl.LightningModule):
    def __init__(self, model, criterion, optimizer):

        super().__init__()
        self.model = model
        self.criterion = criterion
        self.optimizer = optimizer

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        return loss

Lightning Trainer の活用例

重要なポイント:

トレーニングループを管理
分散トレーニングをサポート
コールバックとロギングを処理
リソース使用を最適化

model = LightClassifier()



trainer = Trainer(max_epochs=10, accelerator="gpu", devices=1)
trainer.fit(model, train_dataloader, val_dataloader)

Afro-MNIST データセットの紹介

アフロ・アジア語族とニジェール・コンゴ語族で使用される4つの文字体系（ゲエズ文字（エチオピア文字）、ヴァイ文字、オスマニヤ文字、ンコ文字）に対応する、MNIST 形式の合成データセット群です。

¹ Wu, Daniel J., Andrew C. Yang, and Vinay U. Prabhu. "Afro-MNIST: Synthetic generation of MNIST-style datasets for low-resource languages." arXiv preprint arXiv:2009.13509 (2020).

PyTorch Lightning の振り返り

PyTorch Lightning の図解

練習しましょう！

PyTorch Lightning によるスケーラブルな AI モデルの構築