PyTorch Lightning 入門

Scalable AI Models with PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

PyTorch 與 PyTorch Lightning

$$

標準 PyTorch：

  • 需大量手動處理
  • 手寫訓練迴圈
  • GPU/TPU 管理、記錄與檢查點

程式設計師動圖

Scalable AI Models with PyTorch Lightning

PyTorch 與 PyTorch Lightning

PyTorch Lightning：

  • 建立於 PyTorch 之上
  • 自動化：
    • 訓練
    • 檢查點
    • 記錄
  • 減少樣板程式碼
  • 提升可擴充性與可重現性

PyTorch Lightning 標誌

Scalable AI Models with PyTorch Lightning

PyTorch Lightning 概觀

  • 範例：全球電商優化流程
    • 視覺搜尋模型開發
    • 自動化訓練迴圈
    • 快速迭代、最少樣板

$$

  • 核心元件：LightningModule 與 Trainer
    from lightning.pytorch import LightningModule
from lightning.pytorch import Trainer
Scalable AI Models with PyTorch Lightning

Lightning 結構

關鍵元件：

  • LightningModule：模型核心邏輯

象徵邏輯的燈泡

Scalable AI Models with PyTorch Lightning

Lightning 結構

關鍵元件：

  • LightningModule：模型核心邏輯
  • Lightning Trainer：協調訓練流程

象徵邏輯的燈泡與象徵操作的齒輪

Scalable AI Models with PyTorch Lightning

Lightning 結構

關鍵元件：

  • LightningModule：模型核心邏輯
  • Lightning Trainer：協調訓練流程
  • DataModule：組織資料管線
  • Callbacks：自動化事件
  • Logger：追蹤實驗

象徵邏輯的燈泡與象徵操作的齒輪

Scalable AI Models with PyTorch Lightning

LightningModule 實作示例

重點：

  • __init__：定義模型架構
  • forward()：資料傳遞經模型
  • training_step()：定義訓練
  • 可使用自訂 hook
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
Scalable AI Models with PyTorch Lightning

Lightning Trainer 實作示例

重點：

  • 管理訓練迴圈
  • 支援分散式訓練
  • 處理 callbacks 與記錄
  • 最佳化資源使用
model = LightClassifier()



trainer = Trainer(max_epochs=10, accelerator="gpu", devices=1)
trainer.fit(model, train_dataloader, val_dataloader)
Scalable AI Models with PyTorch Lightning

認識 Afro-MNIST 資料集

一組模擬 MNIST 風格的資料集，涵蓋四種屬於亞非語系與尼日－剛果語系的文字系統：Ge'ez（衣索比亞文）、Vai、Osmanya 與 N'Ko。

資料集封面

1 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).
Scalable AI Models with PyTorch Lightning

PyTorch Lightning 複習

PyTorch Lightning 示意圖

Scalable AI Models with PyTorch Lightning

一起來練習吧！

Scalable AI Models with PyTorch Lightning

Preparing Video For Download...