pytorch-lightning

PyTorch Lightning

概述

PyTorch Lightning 是一个深度学习框架，可整理 PyTorch 代码，减少样板代码，同时保留完全的灵活性。它可以自动化训练工作流、多设备编排，并实现神经网络训练和跨多个 GPU/TPU 扩展的最佳实践。

当前上游版本： lightning 2.6.4（PyPI，2026 年 5 月）。文档：lightning.ai/docs/pytorch/stable。使用 import lightning as L（pytorch-lightning 包名仍会安装同一个库）。

安装

uv pip install lightning

可选附加项：

uv pip install lightning[extra]    # loggers, strategies, etc.
uv pip install wandb mlflow        # specific loggers as needed

何时使用此技能

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

• 使用 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）
• 自动混合精度训练
• 梯度累积和裁剪
• 检查点保存和提前停止
• 进度条和日志记录

快速配置参考： 常见的 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——跟踪 LR 调度器的变化
• BatchSizeFinder——自动确定最佳批次大小

详细文档： 请阅读 references/callbacks.md，了解内置回调和自定义回调的创建方法。

5. 日志记录——实验跟踪

可集成多个日志平台：

• TensorBoard（默认）
• Weights & Biases（WandbLogger）
• MLflow（MLFlowLogger）
• Comet（CometLogger）
• CSV（CSVLogger）

注意：NeptuneLogger 已在 lightning 2.6.4 中移除。请改用 W&B、MLflow 或 TensorBoard。

在任何 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 - Logger 集成和用法
• distributed_training.md - DDP、FSDP、DeepSpeed 比较和设置
• best_practices.md - 常见模式、提示和易错点

Agent 操作流程

1. 检查环境。 确认 Python 环境和库版本（python -c "import pkg; print(pkg.__version__)"），并检查数据的形状、类型和缺失值。
2. 明确输入。 根据数据或用户确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 先在样本或单个折上运行，并检查运行时间和内存占用。
4. 使用以上说明和参考资料执行完整任务。
5. 验证结果。 使用留出数据、固定随机种子和适当的指标；检查是否存在数据泄漏；报告不确定性（置信区间、不同种子下的标准差）。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及哪些方面仍不确定。

如果发生这种情况	这样做
内存不足或执行速度过慢	减少采样，使用分块或惰性计算，或缩小模型规模，并告知用户做了哪些更改。
这些说明中的函数、标志或端点在已安装的版本中不存在	查阅已安装版本自己的文档（help()、--help、官方文档），调整做法并告知用户。切勿编造 API。
必需的输入、标识符或参数含义不明确	询问用户，或在运行前明确说明所采用的假设。

诚信规则

• 切勿捏造结果、参数、标识符、引文或统计数据。如果无法运行或验证某项内容，请如实说明。
• 切勿报告本次会话中未计算的指标；展示生成每个数值的代码路径。
• 此处与特定版本相关的细节可能已经过时：请对照已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，须事先征得同意。

相关技能

• huggingface-accelerate：使用 HuggingFace Accelerate（Accelerator 类、accelerate config、accelerate launch）封装现有 PyTorch 训练脚本，使同一份脚本……
• pytorch-lightning-distributed：高级 PyTorch 框架，提供 Trainer 类、自动分布式训练（DDP/FSDP/DeepSpeed）、回调系统，以及精简的样板代码……
• experiment-tracking-swanlab：使用 SwanLab 跟踪机器学习实验。SwanLab 是一款开源工具，涵盖 swanlab.init、配置和指标记录、标量图表及媒体记录……