当前位置：首页 > news >正文

PyTorch Lightning Trainer介绍

news 2025/8/26 19:04:38

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]