文章网站后台管理系统培训班报名
PyTorch Lightning 的 Trainer 是框架的核心类，负责自动化训练流程、分布式训练、日志记录、模型保存等复杂操作。通过配置参数即可快速实现高效训练，无需手动编写循环代码。以下是详细介绍和使用示例：
Trainer 的核心功能
-
自动化训练循环
自动处理training_step、validation_step、test_step的调用，无需手动编写for epoch in epochs循环。 -
硬件加速支持
支持 CPU/GPU/TPU、多卡训练（DDP、DeepSpeed）、混合精度训练等。 -
训练控制
控制训练轮数 (max_epochs)、批次大小 (batch_size)、梯度裁剪 (gradient_clip_val) 等。 -
日志与监控
集成 TensorBoard、W&B、MLFlow 等日志工具，监控损失、准确率等指标。 -
回调机制
通过回调函数（如EarlyStopping,ModelCheckpoint）实现早停、模型保存等扩展功能。
Trainer 的常用参数
from pytorch_lightning import Trainertrainer = Trainer(# 基础配置max_epochs=10, # 最大训练轮数accelerator="auto", # 自动选择设备 (CPU/GPU/TPU)devices="auto", # 使用所有可用设备（如多 GPU）precision="16-mixed", # 混合精度训练（FP16）# 日志与调试logger=True, # 默认使用 TensorBoardlog_every_n_steps=10, # 每 10 个批次记录一次日志fast_dev_run=False, # 快速运行一个批次（调试模式）# 回调函数callbacks=[pl.callbacks.EarlyStopping(monitor="val_loss", patience=3),pl.callbacks.ModelCheckpoint(monitor="val_loss", save_top_k=2)],# 分布式训练strategy="ddp", # 分布式数据并行策略（多 GPU）num_nodes=1, # 节点数量（多机器训练）
)使用示例代码
步骤 1：定义 LightningModule
import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as plclass LitModel(pl.LightningModule):def __init__(self):super().__init__()self.layer1 = nn.Linear(28*28, 128)self.layer2 = nn.Linear(128, 10)def forward(self, x):x = x.view(x.size(0), -1) # 展平输入x = F.relu(self.layer1(x))x = self.layer2(x)return xdef training_step(self, batch, batch_idx):x, y = batchy_hat = self(x)loss = F.cross_entropy(y_hat, y)self.log("train_loss", loss) # 自动记录日志return lossdef validation_step(self, batch, batch_idx):x, y = batchy_hat = self(x)loss = F.cross_entropy(y_hat, y)self.log("val_loss", loss) # 自动记录验证损失def configure_optimizers(self):return torch.optim.Adam(self.parameters(), lr=0.001)步骤 2：定义 DataModule
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensorclass MNISTDataModule(pl.LightningDataModule):def __init__(self, batch_size=32):super().__init__()self.batch_size = batch_sizedef prepare_data(self):MNIST(root="data", download=True)def setup(self, stage=None):full_dataset = MNIST(root="data", train=True, transform=ToTensor())self.train_data, self.val_data = random_split(full_dataset, [55000, 5000])def train_dataloader(self):return DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True)def val_dataloader(self):return DataLoader(self.val_data, batch_size=self.batch_size)dm = MNISTDataModule(batch_size=32)步骤 3：启动训练
model = LitModel()
trainer = Trainer(max_epochs=10,accelerator="auto",devices="auto",logger=True,callbacks=[pl.callbacks.ModelCheckpoint(monitor="val_loss")]
)# 开始训练与验证
trainer.fit(model, datamodule=dm)# 测试（可选）
trainer.test(model, datamodule=dm)关键功能演示
1. 多 GPU 训练
# 使用 4 个 GPU 训练
trainer = Trainer(devices=4, strategy="ddp")2. 混合精度训练
# 使用 FP16 混合精度
trainer = Trainer(precision="16-mixed")3. 早停与模型保存
callbacks = [pl.callbacks.EarlyStopping(monitor="val_loss", patience=3),pl.callbacks.ModelCheckpoint(dirpath="checkpoints/",filename="best-model-{epoch:02d}-{val_loss:.2f}",save_top_k=2,monitor="val_loss")
]
trainer = Trainer(callbacks=callbacks)4. 调试模式
# 快速验证代码正确性（仅运行一个批次）
trainer = Trainer(fast_dev_run=True)常见问题
如何恢复训练？
使用 resume_from_checkpoint 参数：
trainer = Trainer(resume_from_checkpoint="path/to/checkpoint.ckpt")如何限制训练时间？
trainer = Trainer(max_time="00:02:00") # 最多训练 2 分钟如何自定义学习率调度器？
在 自定义的 LightningDataModule继承类的 configure_optimizers 方法中返回优化器和调度器：
def configure_optimizers(self):optimizer = Adam(self.parameters())scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)return [optimizer], [scheduler]总结
通过 Trainer，PyTorch Lightning 将训练流程的复杂性封装在几行配置中，开发者只需关注模型逻辑和数据加载。其灵活的参数和回调机制能够覆盖从实验到生产的全流程需求。
参考：
https://lightning.ai/docs/pytorch/stable/common/trainer.html