pytorch-lightning

PyTorch Lightning

概述

PyTorch Lightning 是一个深度学习框架，可整理 PyTorch 代码，消除样板代码，同时保持充分的灵活性。它可自动化训练工作流、多设备编排，并帮助实现神经网络训练和扩展到多块 GPU/TPU 的最佳实践。

何时使用此技能

在以下情况下应使用此技能：

• 使用 PyTorch Lightning 构建、训练或部署神经网络
• 将 PyTorch 代码整理到 LightningModules 中
• 配置 Trainer 以进行多 GPU/TPU 训练
• 使用 LightningDataModules 实现数据流水线
• 使用回调、日志记录和分布式训练策略（DDP、FSDP、DeepSpeed）
• 以专业方式组织深度学习项目

核心功能

1. LightningModule — 模型定义

将 PyTorch 模型整理为六个逻辑部分：

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 可自动执行训练循环、设备管理、梯度操作和回调。主要功能：

• 支持多 GPU/TPU，并可选择策略（DDP、FSDP、DeepSpeed）
• 自动混合精度训练
• 梯度累积和裁剪
• 检查点保存和提前停止
• 进度条和日志记录

快速设置参考： 查看 scripts/quick_trainer_setup.py，了解常见的 Trainer 配置。

详细文档： 阅读 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 - 适用于参数量 <500M 的模型（ResNet、较小的 Transformer）
• FSDP - 适用于参数量 500M+ 的模型（大型 Transformer，推荐 Lightning 用户使用）
• DeepSpeed - 适用于前沿功能和精细控制

使用以下方式配置：Trainer(strategy="ddp", accelerator="gpu", devices=4)

详细文档： 阅读 references/distributed_training.md，了解策略比较和配置。

7. 最佳实践

• 与设备无关的代码 - 使用 self.device，而不是 .cuda()
• 超参数保存 - 在 __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 - 常见模式、提示和易错点