Agent Skills
返回列表
💻

PyTorch Lightning

开发编程 更新于 2026.08.29

将以下提示词粘贴到你的 AI 对话框中：

请根据 https://skillhub.cn/install/skillhub.md，安装 @org-02qudk26/pytorch-lightning。

技能介绍

要解决什么问题

直接用 PyTorch 写训练脚本时，for epoch 循环、设备搬运、分布式启动和 checkpoint 往往堆在同一文件里，改一次实验就要复制整段样板。模型变大后还要在 DDP、FSDP、DeepSpeed 之间手动切换，日志和早停也容易和核心逻辑缠在一起。

技能如何组织工作流

本技能围绕 PyTorch Lightning 的三块骨架展开：

  • LightningModule：把模型拆成 training_step / validation_step、优化器配置等固定钩子，而不是自写完整训练循环。
  • Trainer：负责设备、混合精度、梯度累积与裁剪、checkpoint 和进度展示；多卡时通过 strategy（如 ddp、fsdp）选型扩展。
  • LightningDataModule：把 prepare_data、setup 与各阶段 DataLoader 封装成可复用数据管道。

在此基础上，技能还覆盖 回调（如 ModelCheckpoint、EarlyStopping）、self.log() 与多种 Logger（TensorBoard、W&B 等），以及分布式策略的适用场景说明。仓库内 scripts/ 提供 Module、DataModule、Trainer 模板，references/ 对应各组件的详细说明。

适用边界

Lightning 是在 PyTorch 之上做结构约束，不能替代你对 loss、架构和数据正确性的判断。极定制化的训练循环若与 Lightning 钩子模型冲突，可能需要部分逻辑留在 Module 外或评估是否仍适合 Lightning。复杂多步研究型流水线若超出单次 fit 范畴，需自行规划编排方式。

使用场景

  • 现有 PyTorch 训练脚本里循环、设备和保存逻辑缠在一起，需要拆成 LightningModule 并用 Trainer 统一跑通训练与验证。
  • 要在 4 块 GPU 上跑同一套实验，需要配置 Trainer 的 accelerator、devices 以及 DDP 或 FSDP 策略。
  • 数据集下载、划分和 train/val DataLoader 分散在多个脚本里，希望封装成 LightningDataModule 便于复现实验。
  • 训练过程中要自动存最佳权重、指标停滞早停，并把 loss 同步到 TensorBoard 或 W&B。

适合人员

  • 负责维护多条实验分支的算法工程师，希望用 save_hyperparameters 和统一 Trainer 配置减少复制粘贴。
  • 在单机多卡上训练亿级参数 Transformer 的研究员，需要在 DDP 与 FSDP 之间做策略选型并正确配置 Trainer。
  • 把 side project 整理成可交付训练仓库的 ML 工程师，需要 DataModule、回调与日志与模型代码分离。
  • 从裸 PyTorch 迁移到团队规范栈的深度学习开发者，需要对照 Lightning 钩子与最佳实践文档改写法。