pytorch-lightning

PyTorch Lightning - 深度学习训练框架

技能概述

PyTorch Lightning 是一个深度学习框架，用于组织 PyTorch 代码、消除样板代码，同时保持完整的灵活性。它自动化训练工作流、多设备编排，实现神经网络训练和跨多 GPU/TPU 扩展的最佳实践。

适用场景

构建和训练神经网络

当你需要使用 PyTorch 构建深度学习模型，并希望代码结构更加清晰、可维护时。Lightning 帮你将模型定义、训练循环、验证逻辑分离到不同方法中，避免混乱的 for 循环嵌套。

多 GPU 和分布式训练

当模型需要在多个 GPU 或 TPU 上并行训练时，Lightning 提供开箱即用的 DDP、FSDP、DeepSpeed 策略，无需手动处理进程通信和设备管理。

规范化深度学习项目

当你需要组织一个专业的深度学习项目，包括数据管道、训练日志、模型检查点、回调机制等完整功能时，Lightning 提供了标准化的项目结构和最佳实践。

核心功能

LightningModule - 模型组织

将 PyTorch 模型组织成六个逻辑部分：初始化、训练循环、验证循环、测试循环、预测和优化器配置。这种结构让代码更清晰、更易测试和复用。

Trainer - 自动化训练

Trainer 自动处理设备管理、梯度操作、混合精度训练、梯度累积、检查点保存、早停等繁琐工作。只需一行代码即可启用多 GPU 训练。

完整的数据和工具生态

提供 LightningDataModule 封装数据管道、内置常用回调（ModelCheckpoint、EarlyStopping）、支持多种日志平台（TensorBoard、W&B、MLflow）、以及分布式训练策略。

常见问题

PyTorch Lightning 是什么？和原生 PyTorch 有什么区别？

PyTorch Lightning 是建立在 PyTorch 之上的轻量级框架，它不改变 PyTorch 的功能，而是将训练代码组织成更清晰的结构。原生 PyTorch 需要手写训练循环、验证循环、设备管理等大量样板代码，而 Lightning 将这些抽象为 Trainer，让你只需关注模型逻辑和数据处理。

如何开始使用 PyTorch Lightning？

只需将你的 PyTorch 模型继承 LightningModule，实现 training_step 和 configure_optimizers 方法，然后用 Trainer 替换原来的训练循环即可。如果需要多 GPU 训练，只需设置 accelerator="gpu" 和 devices 参数。

DDP、FSDP 和 DeepSpeed 应该选择哪个？

选择取决于模型大小：小于 5 亿参数的模型（如 ResNet、小型 Transformer）推荐 DDP；5 亿参数以上的大型模型推荐 FSDP（Lightning 官方推荐）；如果需要更精细的控制和最新特性，可以选择 DeepSpeed。配置方式都是 Trainer(strategy="...")。