mcpField
A
alterlab-pytorch-lightning

alterlab-pytorch-lightning

github.com/alterlab-ieu数据与数据库7866 次安装5 周前更新
代码仓库

关于 alterlab-pytorch-lightning

alterlab-pytorch-lightning 是使用 PyTorch Lightning 组织和训练神经网络的技能，介绍可复用的模型与数据模块、训练自动化、回调、日志记录及多设备策略。

主要功能

  1. 01LightningModule 模型组织
  2. 02Trainer 训练流程自动化
  3. 03可复用的 LightningDataModule 数据管线
  4. 04用于检查点和早停的回调
  5. 05日志记录与分布式训练策略

适用场景

  1. 01组织可复用的神经网络训练项目
  2. 02在 GPU 或 TPU 上运行训练
  3. 03添加检查点、早停和指标日志
  4. 04配置 DDP、FSDP 或 DeepSpeed

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/alterlab-ieu/alterlab-academic-skills --skill alterlab-pytorch-lightning

安装后在 Claude Code 中输入/alterlab-pytorch-lightning确认已加载。

包含的文件

除 SKILL.md 外，安装时一起下载的参考资料和脚本。

evals/
evals.json
references/
best_practices.md
callbacks.md
data_module.md
distributed_training.md
lightning_module.md
logging.md
trainer.md
scripts/
quick_trainer_setup.py
template_datamodule.py
template_lightning_module.py

详细介绍

PyTorch Lightning

概述

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

uv pip install lightning   # provides `import lightning as L` and `lightning.pytorch.*`

本技能中的所有代码都使用统一的 lightning 包。单独发布的 pytorch-lightning 包安装的是较旧的 pytorch_lightning 命名空间；请勿在同一个项目中混用这两个命名空间，因为它们的类彼此独立（例如，lightning.Trainer.fit 会因 TypeError 拒绝 pytorch_lightning.LightningModule）。

何时使用此技能

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

  • 使用 PyTorch Lightning 构建、训练或部署神经网络
  • 将 PyTorch 代码整理为 LightningModules
  • 配置 Trainer 以进行多 GPU/TPU 训练
  • 使用 LightningDataModules 实现数据管道
  • 使用回调、日志记录和分布式训练策略（DDP、FSDP、DeepSpeed）
  • 以专业方式组织深度学习项目
不适用的情况
场景改用
使用 transformers Trainer / pipelines 微调预训练的 Hugging Face 模型alterlab-transformers
设计图神经网络层或加载图数据集（GCN、GAT、消息传递）alterlab-torch-geometric
在环境中训练强化学习 Agent（PPO、SAC、DQN）alterlab-stable-baselines3
对表格数据进行传统机器学习（随机森林、梯度提升、线性模型）alterlab-scikit-learn

核心能力

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 — 跟踪 LR 调度器的变化
  • BatchSizeFinder — 自动确定最佳批量大小

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

5. 日志记录 — 实验跟踪

集成多个日志平台：

  • TensorBoard（安装了 tensorboard 时默认使用；否则使用 CSV）
  • Weights & Biases（WandbLogger）
  • MLflow（MLFlowLogger）
  • Comet（CometLogger）
  • CSV（CSVLogger）

Neptune 的托管服务关闭后（2026 年 3 月），NeptuneLogger 已在 Lightning 2.6.4 中移除。

在任何 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 — 常见模式、技巧和注意事项

评论

支持 Markdown