PyTorch Lightning 简介

使用 PyTorch Lightning 构建可扩展 AI 模型

Sergiy Tkachuk

Director, GenAI Productivity

PyTorch 与 PyTorch Lightning

$$

标准 PyTorch：

  • 需大量手工工作
  • 需显式编写训练循环
  • 需处理 GPU/TPU、日志与检查点

coder.gif

使用 PyTorch Lightning 构建可扩展 AI 模型

PyTorch 与 PyTorch Lightning

PyTorch Lightning：

  • 构建于 PyTorch 之上
  • 自动化：
    • 训练
    • 检查点
    • 日志
  • 减少样板代码
  • 提升可扩展性与可复现性

PyTorch Lightning 标志

使用 PyTorch Lightning 构建可扩展 AI 模型

PyTorch Lightning 概览

  • 示例：全球电商精简流程
    • 视觉搜索模型开发
    • 训练循环自动化
    • 以最少样板代码快速迭代

$$

  • 核心组件：LightningModule 与 Trainer
    from lightning.pytorch import LightningModule
from lightning.pytorch import Trainer
使用 PyTorch Lightning 构建可扩展 AI 模型

Lightning 结构

关键组件：

  • LightningModule：模型核心逻辑

一个代表逻辑的灯泡

使用 PyTorch Lightning 构建可扩展 AI 模型

Lightning 结构

关键组件：

  • LightningModule：模型核心逻辑
  • Lightning Trainer：统筹训练

一个代表逻辑的灯泡和一个代表操作的齿轮

使用 PyTorch Lightning 构建可扩展 AI 模型

Lightning 结构

关键组件：

  • LightningModule：模型核心逻辑
  • Lightning Trainer：统筹训练
  • DataModule：组织数据管道
  • Callbacks：自动化事件
  • Logger：记录实验

一个代表逻辑的灯泡和一个代表操作的齿轮

使用 PyTorch Lightning 构建可扩展 AI 模型

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
使用 PyTorch Lightning 构建可扩展 AI 模型

Lightning Trainer 实战

要点：

  • 管理训练循环
  • 支持分布式训练
  • 处理回调与日志
  • 优化资源使用
model = LightClassifier()



trainer = Trainer(max_epochs=10, accelerator="gpu", devices=1)
trainer.fit(model, train_dataloader, val_dataloader)
使用 PyTorch Lightning 构建可扩展 AI 模型

Afro-MNIST 数据集介绍

一组用于阿非罗-亚细亚语系与尼日尔-刚果语系中四种文字的合成 MNIST 风格数据集：吉兹文（埃塞俄比亚文字）、瓦伊文、奥斯曼亚文与恩科文。

dataset-cover.png

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).
使用 PyTorch Lightning 构建可扩展 AI 模型

PyTorch Lightning 回顾

PyTorch Lightning 示意图

使用 PyTorch Lightning 构建可扩展 AI 模型

Vamos praticar!

使用 PyTorch Lightning 构建可扩展 AI 模型

Preparing Video For Download...