pytorch-lightning

# PyTorch Lightning ## 概述 PyTorch Lightning 是一个深度学习框架，它通过组织 PyTorch 代码来消除样板代码，同时保持完全的灵活性。它可以自动化训练工作流、多设备编排，并实现神经网络训练和跨多个 GPU/TPU 扩展的最佳实践。 ## 何时使用此技能 当出现以下情况时，应使用此技能： - 使用 PyTorch Lightning 构建、训练或部署神经网络 - 将 PyTorch 代码组织为 LightningModules - 为多 GPU/TPU 训练配置训练器 (Trainer) - 使用 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` 获取全面的方法文档、钩子 (hooks)、属性和最佳实践。 ### 2. Trainer - 训练自动化 训练器 (Trainer) 可自动化训练循环、设备管理、梯度操作和回调。主要功能： - 支持多 GPU/TPU，并可选择策略 (DDP, FSDP, DeepSpeed) - 自动混合精度训练 - 梯度累积和裁剪 - 检查点保存和提前停止 - 进度条和日志记录 **快速设置参考：** 请参阅 `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` 获取一个 co