PyTorch Lightning实战：从手写训练循环到标准化训练流程

1. 从手写训练循环到Lightning的思维转换

如果你写过一段时间的PyTorch，大概率经历过这样的场景：模型定义写完了，兴冲冲准备训练，结果发现要手写训练循环、验证循环、指标计算、日志记录、模型保存、学习率调度、混合精度、多卡并行……一套下来，模型本身可能只占代码量的三分之一，剩下全是工程脚手架。更麻烦的是，每次换一个项目，这些脚手架代码就得重新抄一遍、改一遍，稍不留神就漏掉某个关键步骤，比如忘了在验证时调用model.eval()，或者忘了torch.no_grad()，导致显存爆炸或者指标异常。

PyTorch Lightning（后面简称PL）要解决的就是这个问题。它把训练过程中那些“跟模型本身无关、但每次都得写”的工程代码抽离出来，封装成标准化的流程，让你只需要关注模型结构、损失函数、优化器配置和数据处理这几件事。你可以把它理解成给PyTorch加了一层“训练流程模板”，你负责填内容，它负责跑流程。

我第一次接触PL的时候，心里其实是有点抵触的——又多学一套框架，值不值得？但真正用下来之后发现，它并没有隐藏PyTorch的细节，你依然可以用原生的torch.nn.Module、torch.optim、DataLoader，只是在外面包了一层LightningModule和Trainer。想深入底层的时候随时可以钻进去，想快速跑实验的时候它又能帮你省掉大量重复劳动。

这篇文章面向的是已经会写基础PyTorch代码、但还没系统用过PL的开发者。我会从实际使用角度出发，把PL的核心概念、标准使用步骤、常见坑点和进阶技巧都讲清楚，让你看完之后能直接上手跑自己的项目。

2. LightningModule的骨架：把训练逻辑拆成标准零件

2.1 为什么是LightningModule而不是普通nn.Module

PL的核心抽象是LightningModule，它继承自torch.nn.Module，所以你的模型定义方式完全不变。区别在于，你可以在同一个类里面把训练相关的逻辑都组织起来，包括：

• __init__：定义网络层结构
• forward：定义前向传播
• training_step：定义单步训练做什么
• validation_step：定义单步验证做什么
• configure_optimizers：定义优化器和学习率调度器
• train_dataloader/val_dataloader：定义数据加载

这种组织方式的好处是，所有跟这个模型相关的逻辑都集中在一个类里面，而不是散落在训练脚本的各个角落。当你需要复现实验或者迁移到另一个项目时，直接把这个类复制过去就行，不用再去翻训练脚本里那些零散的配置。

我个人的体会是，这种“把训练逻辑内聚到模型类”的做法，在项目初期可能感觉多写了几行代码，但到了中期需要频繁调整训练策略的时候，优势就非常明显了——你只需要改一个地方，而不是在多个文件之间来回跳。

2.2 training_step里到底该放什么

training_step是PL训练流程的核心钩子，每一步训练都会调用它。它的签名通常是：

def training_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = self.loss_fn(logits, y) self.log("train_loss", loss) return loss

这里有几个关键点需要注意。第一，batch是DataLoader返回的一个批次数据，batch_idx是当前批次的索引。第二，你必须返回loss，因为PL会根据这个loss自动进行反向传播和参数更新。第三，self.log是PL提供的日志接口，它会自动处理分布式训练中的日志聚合问题。

很多人刚开始用的时候会习惯性地在training_step里面手动调用loss.backward()和optimizer.step()，这是不需要的，PL会自动帮你做。你只需要把前向传播和loss计算写好就行。

注意：training_step里面不要做跟当前批次无关的操作，比如修改学习率、保存模型等，这些应该放在对应的钩子函数里。

2.3 configure_optimizers的返回值有讲究

configure_optimizers负责告诉PL用什么优化器和调度器。最简单的写法是只返回一个优化器：

def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=1e-3)

但如果要用学习率调度器，就需要返回一个字典：

def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "interval": "epoch", "frequency": 1, }, }

这里的interval可以是"step"或"epoch"，决定了调度器是按步更新还是按轮更新。frequency表示每隔多少个interval调用一次。这两个参数如果设错了，学习率的变化节奏就会跟你预期的不一样，训练效果可能差很多。

我踩过的一个坑是：用ReduceLROnPlateau的时候，需要把monitor参数设成对应的指标名，否则调度器不知道根据哪个指标来调整学习率。而且这个指标必须是通过self.log记录过的，否则会报错找不到。

3. Trainer的配置：那些影响训练效率和稳定性的参数

3.1 max_epochs和max_steps的区别与选择

Trainer是PL的训练引擎，你通过配置它来控制训练过程。最基础的两个参数是max_epochs和max_steps：

• max_epochs：训练多少个完整轮次
• max_steps：训练多少个批次步

这两个参数只需要设一个，如果都设了，以先达到的为准。对于大多数中小规模数据集，用max_epochs更直观；对于大规模数据集或者需要精确控制训练量的场景，用max_steps更方便。

我一般会在调试阶段用max_steps设一个很小的值（比如100），快速验证整个流程能不能跑通，确认没问题之后再换成max_epochs做完整训练。

3.2 accelerator和devices的正确配置方式

PL支持多种硬件后端，通过accelerator和devices两个参数来控制：

# 单GPU trainer = Trainer(accelerator="gpu", devices=1) # 多GPU trainer = Trainer(accelerator="gpu", devices=4) # 自动选择 trainer = Trainer(accelerator="auto", devices="auto")

accelerator可以设为"cpu"、"gpu"、"auto"等。devices可以是一个整数（表示用几张卡），也可以是一个列表（指定用哪几张卡）。

这里有个实际经验：如果你在本地开发机上调试，建议显式设成accelerator="cpu"，避免不小心占用了GPU资源。到了服务器上再改成GPU配置。另外，多卡训练的时候，PL默认用的是DDP（DistributedDataParallel）策略，比DP（DataParallel）效率高很多，但要求你的代码里不能有跨卡不兼容的操作。

3.3 梯度累积与混合精度的配合使用

当显存不够但又想用大batch size的时候，梯度累积是一个常用技巧。PL通过accumulate_grad_batches参数支持这个功能：

trainer = Trainer( accumulate_grad_batches=4, precision="16-mixed", )

上面这个配置表示：每4个批次才做一次参数更新，等效于把batch size扩大了4倍。同时开启了混合精度训练，用FP16做前向和反向计算，用FP32做参数更新，既能省显存又能加速。

但要注意，梯度累积和BatchNorm层一起用的时候可能会有问题，因为BN统计量是按实际批次计算的，累积之后的等效batch size跟BN看到的batch size不一致。如果发现训练不稳定，可以尝试把BN换成GroupNorm或者LayerNorm。

4. 数据加载与回调：让训练流程更顺滑的细节

4.1 DataLoader的配置要点

PL对DataLoader没有特殊要求，你依然用原生的torch.utils.data.DataLoader。但有几个参数值得注意：

from torch.utils.data import DataLoader train_loader = DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, )

num_workers设成CPU核心数的一半到相等之间比较合适，太少会导致数据加载成为瓶颈，太多会浪费内存。pin_memory=True在GPU训练时能加速数据传输。drop_last=True在梯度累积时建议开启，避免最后一个不完整的批次影响累积效果。

在LightningModule里，你通过train_dataloader方法返回这个loader：

def train_dataloader(self): return train_loader

PL会自动在每轮训练开始前调用这个方法获取数据加载器。

4.2 ModelCheckpoint：保存最佳模型的正确姿势

ModelCheckpoint是PL最常用的回调之一，用来在训练过程中保存模型权重：

from lightning.pytorch.callbacks import ModelCheckpoint checkpoint_callback = ModelCheckpoint( dirpath="checkpoints/", filename="best-{epoch}-{val_loss:.4f}", monitor="val_loss", mode="min", save_top_k=3, save_last=True, )

monitor指定要监控的指标，mode指定是越小越好还是越大越好，save_top_k表示保留最好的前K个模型，save_last表示额外保存最后一个epoch的模型。

这里有个容易忽略的点：monitor指定的指标必须是在validation_step里通过self.log记录过的，而且self.log的on_epoch参数要设为True（默认就是True），否则PL拿不到这个指标的值。

4.3 EarlyStopping：什么时候该停就停

EarlyStopping用来在验证指标不再提升时提前终止训练，避免过拟合和浪费计算资源：

from lightning.pytorch.callbacks import EarlyStopping early_stop_callback = EarlyStopping( monitor="val_loss", patience=5, mode="min", verbose=True, )

patience表示容忍多少个epoch没有提升，超过这个数就停止训练。这个值设得太小容易过早停止，设得太大又失去了早停的意义。我的经验是，对于大多数任务，patience设在5到10之间比较合理。如果数据集比较小、指标波动大，可以适当调大。

4.4 LearningRateMonitor：观察学习率变化

这个回调会记录学习率的变化并输出到日志中，方便你判断调度器是否按预期工作：

from lightning.pytorch.callbacks import LearningRateMonitor lr_monitor = LearningRateMonitor(logging_interval="epoch")

把它加到Trainer的callbacks列表里就行。训练结束后，你可以在TensorBoard或者WandB里看到学习率曲线，确认它是否按照你设定的策略在变化。

5. 从单卡到多卡：PL帮你处理了哪些麻烦事

5.1 DDP策略下需要注意的代码写法

PL在多卡训练时默认使用DDP策略，每个进程有独立的模型副本和数据加载器。这意味着你的代码需要满足一些额外要求：

• 不要在__init__里做跟设备相关的操作，因为每个进程都会执行一次
• 不要手动调用.cuda()或.to(device)，PL会自动处理
• 日志记录用self.log而不是print，PL会自动聚合各进程的日志
• 如果有需要跨进程同步的操作，用self.all_gather等PL提供的接口

我遇到过一个典型问题：在training_step里用print输出调试信息，结果多卡训练时每个进程都打印一遍，日志刷屏。后来改成用self.log记录到TensorBoard，清爽多了。

5.2 分布式采样器的自动处理

在DDP模式下，每个进程需要看到不同的数据子集，PL会自动为DataLoader配置DistributedSampler。你不需要手动设置，但要注意shuffle参数的行为会有所变化——PL会接管shuffle逻辑，你设的shuffle=True依然有效，但实际的分片由采样器控制。

如果你的数据集有特殊的分片需求（比如按类别分层采样），可以通过train_dataloader返回自定义的采样器，PL会尊重你的设置。

5.3 多卡训练时的日志聚合

PL的self.log方法有一个sync_dist参数，默认是False。在多卡训练时，如果你希望记录的指标是所有卡上的平均值，需要设成True：

self.log("val_loss", loss, sync_dist=True)

如果不设，每个进程会记录自己卡上的值，最终日志里看到的是各卡结果的拼接，而不是全局平均。对于验证指标来说，通常需要全局平均才有意义，所以建议在validation_step里记录指标时都加上sync_dist=True。

6. 那些文档里不会写的实战经验

6.1 调试阶段先用fast_dev_run

PL提供了一个非常实用的调试参数fast_dev_run：

trainer = Trainer(fast_dev_run=True)

开启后，PL只会跑一个批次的训练和一个批次的验证，然后结束。这能在几秒钟内帮你验证整个流程是否有语法错误、维度不匹配、设备配置错误等问题。我养成的习惯是，每次写完一个新的LightningModule，先用fast_dev_run=True跑一遍，确认没问题再开始正式训练。

6.2 用seed_everything保证可复现性

PL提供了seed_everything函数来统一设置随机种子：

from lightning.pytorch import seed_everything seed_everything(42, workers=True)

workers=True会确保DataLoader的worker进程也使用相同的种子。这个函数应该在创建Trainer之前调用。但要注意，即使设了种子，由于GPU计算的非确定性，完全复现结果仍然需要设置torch.use_deterministic_algorithms(True)，但这可能会降低训练速度。

6.3 模型保存与加载的常见误区

PL保存的checkpoint包含模型权重、优化器状态、调度器状态、当前epoch等信息。加载的时候有两种方式：

# 方式一：加载完整checkpoint继续训练 model = MyLightningModule.load_from_checkpoint("checkpoints/best.ckpt") # 方式二：只加载权重用于推理 model = MyLightningModule() checkpoint = torch.load("checkpoints/best.ckpt") model.load_state_dict(checkpoint["state_dict"])

方式一适合继续训练，方式二适合部署推理。需要注意的是，load_from_checkpoint会调用__init__，所以如果你的__init__需要参数，必须通过load_from_checkpoint的关键字参数传进去，否则会报错。

6.4 自定义Callback的典型场景

当PL内置的回调不能满足需求时，你可以自定义Callback：

from lightning.pytorch.callbacks import Callback class MyCallback(Callback): def on_train_epoch_end(self, trainer, pl_module): # 每个训练epoch结束时执行 pass def on_validation_epoch_end(self, trainer, pl_module): # 每个验证epoch结束时执行 pass

常见的自定义场景包括：动态调整某个超参数、在特定epoch保存额外信息、向外部系统发送训练状态等。自定义Callback的好处是逻辑独立，不会污染LightningModule的代码。

6.5 常见报错与排查思路

排查PL的问题时，一个有效的方法是先用fast_dev_run=True和accelerator="cpu"跑一遍，排除硬件和分布式相关的因素，把问题范围缩小到代码逻辑本身。

7. 一个完整的训练脚本长什么样

把前面讲的内容串起来，一个典型的PL训练脚本大概是这样的结构：

import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset from lightning.pytorch import LightningModule, Trainer, seed_everything from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping class MyModel(LightningModule): def __init__(self, input_dim=10, hidden_dim=64, lr=1e-3): super().__init__() self.save_hyperparameters() self.net = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1), ) self.loss_fn = nn.MSELoss() def forward(self, x): return self.net(x) def training_step(self, batch, batch_idx): x, y = batch pred = self(x) loss = self.loss_fn(pred, y) self.log("train_loss", loss, prog_bar=True) return loss def validation_step(self, batch, batch_idx): x, y = batch pred = self(x) loss = self.loss_fn(pred, y) self.log("val_loss", loss, prog_bar=True, sync_dist=True) return loss def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=self.hparams.lr) def train_dataloader(self): x = torch.randn(1000, 10) y = torch.randn(1000, 1) return DataLoader(TensorDataset(x, y), batch_size=32, shuffle=True) def val_dataloader(self): x = torch.randn(200, 10) y = torch.randn(200, 1) return DataLoader(TensorDataset(x, y), batch_size=32) if __name__ == "__main__": seed_everything(42, workers=True) model = MyModel() checkpoint = ModelCheckpoint( dirpath="checkpoints/", filename="best-{epoch}-{val_loss:.4f}", monitor="val_loss", mode="min", save_top_k=1, ) early_stop = EarlyStopping(monitor="val_loss", patience=5, mode="min") trainer = Trainer( max_epochs=50, accelerator="auto", devices="auto", callbacks=[checkpoint, early_stop], log_every_n_steps=10, ) trainer.fit(model)

这个脚本可以直接运行，你只需要把数据加载部分换成自己的数据集就行。save_hyperparameters()会自动把__init__的参数保存到checkpoint里，加载模型时不用再手动传一遍。

8. 关于PL值不值得用的个人看法

用了两年多PL之后，我的感受是：它适合那些需要频繁做实验、经常调整训练策略、或者需要在多卡环境下跑训练的场景。如果你的项目非常简单，就是一个固定的训练流程跑到底，那用原生PyTorch写一个脚本可能更直接。但只要你开始遇到“这个实验要改一下学习率策略”“那个实验要加个早停”“这次要上多卡”这类需求，PL的标准化流程就能帮你省下大量时间。

它最大的价值不是让你少写代码，而是让你的训练代码有一个统一的结构。当你看别人的PL代码时，你知道training_step里一定是训练逻辑，configure_optimizers里一定是优化器配置，不用去猜每个脚本的组织方式。这种一致性在团队协作和代码复用时特别有价值。

当然，PL也不是没有缺点。它的抽象层在某些边缘场景下会显得不够灵活，比如你想实现一个非常规的训练流程，可能需要绕过一些默认行为。但好在它并没有把你锁死，你随时可以回到原生PyTorch的写法，PL只是提供了一层可选的便利。