pytorch-lightning

# PyTorch Lightning - 高级训练框架 ## 快速开始 PyTorch Lightning 组织 PyTorch 代码以消除样板代码，同时保持灵活性。 **安装**: bash pip install lightning **将 PyTorch 转换为 Lightning** (3 个步骤): python import lightning as L import torch from torch import nn from torch.utils.data import DataLoader, Dataset # 第 1 步：定义 LightningModule (组织你的 PyTorch 代码) class LitModel(L.LightningModule): def __init__(self, hidden_size=128): super().__init__() self.model = nn.Sequential( nn.Linear(28 * 28, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 10) ) def training_step(self, batch, batch_idx): x, y = batch y_hat = self.model(x) loss = nn.functional.cross_entropy(y_hat, y) self.log('train_loss', loss) # 自动记录到 TensorBoard return loss def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=1e-3) # 第 2 步：创建数据 train_loader = DataLoader(train_dataset, batch_size=32) # 第 3 步：使用 Trainer 进行训练 (处理其他所有事情!) trainer = L.Trainer(max_epochs=10, accelerator='gpu', devices=2) model = LitModel() trainer.fit(model, train_loader) **就是这样！** Trainer 处理以下内容： - GPU/TPU/CPU 切换 - 分布式训练 (DDP, FSDP, DeepSpeed) - 混合精度 (FP16, BF16) - 梯度累积 - 检查点保存 - 日志记录 - 进度条 ## 常见工作流 ### 工作流 1：从 PyTorch 到 Lightning **原始 PyTorch 代码**: python model = MyModel() optimizer = torch.optim.Adam(model.parameters()) model.to('cuda') for epoch in range(max_epochs): for batch in train_loader: batch = batch.to('cuda') optimizer.zero_grad() loss = model(batch) loss.backward() optimizer.step() **Lightning 版本**: python class LitModel(L.LightningModule): def __init__(self): super().__init__() self.model = MyModel() def training_step(self, batch, batch_idx): loss = self.model(batch) # 不需要 .to('cuda')! return loss def configure_optimizers(self): return torch.optim.Adam(self.parameters()) # 训练 trainer = L.Trainer(max_epochs=10, accelerator='gpu') trainer.fit(LitModel(), train_loader) **优势**: 40+ 行代码 → 15 行，无需设备管理，自动分布式 ### 工作流 2：验证与测试 python class LitMo