PyTorch Lightning 完全指南：让深度学习训练更简洁高效

前言：为什么选择 PyTorch Lightning？

在深度学习领域，PyTorch 以其灵活性和强大的功能成为研究者和工程师的首选框架。然而，随着项目规模的增长，vanilla PyTorch 代码中充斥着大量重复的训练循环、样板代码和设备管理逻辑。这不仅降低了代码可读性，还增加了维护成本。

PyTorch Lightning 应运而生——它是 PyTorch 的轻量级封装，旨在将研究代码与工程代码分离，让你专注于模型架构本身，而将训练、验证、日志记录、分布式训练等工程细节交给框架自动处理。

Lightning 的核心优势

1. 消除样板代码

Vanilla PyTorch 训练循环通常需要 50+ 行代码，而 Lightning 将其减少到零——你只需定义模型逻辑。

# Vanilla PyTorch - 需要手动管理训练循环
for epoch in range(epochs):
    for batch in train_loader:
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()
    # 验证循环、日志记录、检查点保存...

# PyTorch Lightning - 只需定义模型逻辑
class LitModel(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        loss = self(batch)
        return loss
# Trainer 自动处理所有训练逻辑

2. 无缝扩展到多 GPU/多节点

一行代码即可从 CPU 训练扩展到多 GPU、多节点，甚至 TPU：

# CPU
trainer = Trainer(max_epochs=10)

# 单 GPU
trainer = Trainer(max_epochs=10, accelerator="gpu", devices=1)

# 多 GPU (DataParallel / DistributedDataParallel)
trainer = Trainer(max_epochs=10, accelerator="gpu", devices=4, strategy="ddp")

# 多节点集群
trainer = Trainer(max_epochs=10, accelerator="gpu", devices=8, num_nodes=4)

3. 内置最佳实践

自动混合精度训练 (AMP)：提升训练速度，减少显存占用
梯度累积：模拟大 batch size
学习率监控：自动检测学习率异常
模型检查点：自动保存最佳模型
早期停止：防止过拟合
40+ Callbacks：扩展训练流程

4. 标准化的实验管理

自动记录所有超参数
集成 TensorBoard、W&B、MLflow、Comet 等日志工具
实验可复现性保证
模型版本管理

5. 最新特性（2025）

PyTorch Lightning 2.5+ 引入了多项新特性：

分布式张量支持：更高效的大模型训练
TorchAO 集成：更快的训练和推理优化
改进的编译支持：与 torch.compile 更好地集成
性能提升：减少开销，提升训练速度

第一部分：快速入门

1.1 安装 PyTorch Lightning

# 标准安装
pip install pytorch-lightning

# 安装特定版本
pip install pytorch-lightning==2.5.0

# 安装额外功能（推荐）
pip install pytorch-lightning[extra]  # 包含所有可选依赖

# 从源码安装（开发者）
git clone https://github.com/Lightning-AI/lightning.git
cd lightning
pip install .

验证安装：

import pytorch_lightning as pl
print(pl.__version__)  # 输出版本号，如 2.5.0

1.2 第一个 Lightning 模型

让我们从一个简单的 MNIST 分类器开始，理解 Lightning 的核心模式。

完整示例代码

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

# 1. 定义 LightningModule
class MNISTModel(pl.LightningModule):
    def __init__(self, hidden_dim=128, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()  # 保存超参数

        # 模型层
        self.layer_1 = nn.Linear(28 * 28, hidden_dim)
        self.layer_2 = nn.Linear(hidden_dim, 10)
        self.val_accuracy = []
        self.test_accuracy = []

    def forward(self, x):
        # 前向传播
        x = x.view(x.size(0), -1)  # 展平: [batch, 1, 28, 28] -> [batch, 784]
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        return x

    def training_step(self, batch, batch_idx):
        # 训练步骤
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("train_loss", loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        # 验证步骤
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        preds = torch.argmax(y_hat, dim=1)
        acc = (preds == y).float().mean()
        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)
        return {"val_loss": loss, "val_acc": acc}

    def test_step(self, batch, batch_idx):
        # 测试步骤
        x, y = batch
        y_hat = self(x)
        preds = torch.argmax(y_hat, dim=1)
        acc = (preds == y).float().mean()
        self.log("test_acc", acc)
        return acc

    def configure_optimizers(self):
        # 配置优化器
        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

    def prepare_data(self):
        # 下载/准备数据（只在主进程执行一次）
        datasets.MNIST(".", train=True, download=True)
        datasets.MNIST(".", train=False, download=True)

    def setup(self, stage=None):
        # 设置数据集（在每个进程上执行）
        transform = transforms.Compose([transforms.ToTensor()])
        if stage == "fit" or stage is None:
            mnist_full = datasets.MNIST(".", train=True, transform=transform)
            self.mnist_train, self.mnist_val = random_split(
                mnist_full, [55000, 5000]
            )
        if stage == "test" or stage is None:
            self.mnist_test = datasets.MNIST(".", train=False, transform=transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=32, num_workers=4)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=32, num_workers=4)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=32, num_workers=4)


# 2. 训练模型
if __name__ == "__main__":
    model = MNISTModel()

    # 创建 Trainer
    trainer = pl.Trainer(
        max_epochs=5,
        accelerator="auto",  # 自动检测可用硬件
        devices="auto",      # 自动使用所有可用设备
        precision=16,        # 混合精度训练
    )

    # 训练
    trainer.fit(model)

    # 测试
    trainer.test(model)

    # 预测
    # predictions = trainer.predict(model, test_dataloader)

代码解析：

LightningModule：继承自 pl.LightningModule，封装模型逻辑
training_step：定义单个训练步骤
configure_optimizers：配置优化器
Trainer：管理整个训练流程
数据准备：prepare_data + setup + *_dataloader 模式

第二部分：核心概念深度解析

2.1 LightningModule 详解

LightningModule 是 Lightning 的核心，它将模型逻辑、训练逻辑和数据处理逻辑组织在一起。

核心方法

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        # 1. 初始化模型层
        self.model = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        # 2. 前向传播（用于推理）
        return self.model(x)

    def training_step(self, batch, batch_idx):
        # 3. 训练步骤
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)

        # 记录指标
        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)

        # 可选：返回任意值（会被 callbacks 使用）
        return {"loss": loss, "predictions": y_hat}

    def validation_step(self, batch, batch_idx):
        # 4. 验证步骤
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("val_loss", loss)

    def test_step(self, batch, batch_idx):
        # 5. 测试步骤
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("test_loss", loss)

    def predict_step(self, batch, batch_idx):
        # 6. 预测步骤
        x, _ = batch
        return self(x)

    def configure_optimizers(self):
        # 7. 配置优化器和学习率调度器
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=10
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val_loss",  # 基于验证损失调整学习率
            }
        }

日志记录详解

self.log() 方法提供了丰富的日志记录选项：

# 基础用法
self.log("loss", loss)

# 高级选项
self.log(
    "train_loss",
    loss,
    on_step=True,        # 每个步骤记录
    on_epoch=True,       # 每个 epoch 记录（自动平均）
    prog_bar=True,       # 在进度条显示
    logger=True,         # 记录到日志器
    sync_dist=True,      # 分布式训练时同步
    enable_graph=False,  # 不计算梯度（节省内存）
)

# 记录多个指标
self.log_dict({
    "train_loss": loss,
    "train_acc": accuracy,
    "train_f1": f1_score
}, prog_bar=True)

训练生命周期钩子

Lightning 提供了丰富的钩子函数，可以在训练的不同阶段插入自定义逻辑：

class LitModel(pl.LightningModule):
    def on_train_start(self):
        """训练开始前"""
        print("Starting training...")

    def on_train_epoch_start(self):
        """每个 epoch 开始前"""
        print(f"Epoch {self.current_epoch} starting...")

    def on_train_batch_start(self, batch, batch_idx):
        """每个 batch 开始前"""
        if batch_idx == 0:
            print(f"First batch of epoch {self.current_epoch}")

    def on_train_batch_end(self, outputs, batch, batch_idx):
        """每个 batch 结束后"""
        pass

    def on_train_epoch_end(self):
        """每个 epoch 结束后"""
        # 示例：记录当前学习率
        opt = self.optimizers()
        current_lr = opt.param_groups[0]["lr"]
        self.log("lr", current_lr)

    def on_train_end(self):
        """训练结束后"""
        print("Training completed!")

2.2 Trainer 完全指南

Trainer 是 Lightning 的引擎，负责管理整个训练流程。

基础配置

from pytorch_lightning import Trainer

trainer = Trainer(
    # ===== 训练控制 =====
    max_epochs=100,           # 最大 epoch 数
    max_steps=10000,          # 最大步数（与 max_epochs 二选一）
    max_time="00:12:00:00",   # 最大训练时间（DD:HH:MM:SS）

    # ===== 硬件加速 =====
    accelerator="auto",       # "cpu", "gpu", "tpu", "mps", "auto"
    devices="auto",           # 设备数量：1, 4, "auto"
    precision=16,             # 32 (全精度), 16 (混合精度), "bf16"
    amp_backend="native",     # "native" (apex), "native"

    # ===== 批处理 =====
    accumulate_grad_batches=4,  # 梯度累积步数

    # ===== 检查点 =====
    enable_checkpointing=True,
    default_root_dir="checkpoints/",  # 检查点保存目录

    # ===== 日志 =====
    logger=True,              # 自动使用 TensorBoard
    log_every_n_steps=50,     # 日志记录频率

    # ===== 验证 =====
    check_val_every_n_epoch=1,     # 每个 epoch 验证一次
    val_check_interval=0.5,        # 每个 epoch 一半时验证
    limit_val_batches=1.0,         # 验证时使用的批次比例

    # ===== 早期停止 =====
    enable_progress_bar=True,
    enable_model_summary=True,

    # ===== 确定性 =====
    deterministic=True,        # 保证可复现性
    benchmark=False,           # 使用 cuDNN 自动调优（更快但不可复现）
)

完整配置示例

trainer = Trainer(
    # 多 GPU 配置
    accelerator="gpu",
    devices=4,
    strategy="ddp",  # "dp", "ddp", "ddp_spawn", "fsdp"

    # 混合精度
    precision=16,

    # 训练控制
    max_epochs=50,
    gradient_clip_val=1.0,           # 梯度裁剪
    gradient_clip_algorithm="norm",  # "norm" or "value"

    # 检查点
    callbacks=[
        ModelCheckpoint(
            dirpath="checkpoints/",
            filename="model-{epoch:02d}-{val_loss:.2f}",
            save_top_k=3,
            monitor="val_loss",
            mode="min",
            save_last=True,
        ),
        EarlyStopping(
            monitor="val_loss",
            patience=5,
            mode="min",
        ),
    ],

    # 日志
    logger=TensorBoardLogger("logs/", name="my_model"),

    # 进度条
    enable_progress_bar=True,
    progress_bar_refresh_rate=20,

    # 其他
    weights_summary="full",  # "top", "full", None
    profiler="simple",        # "simple", "advanced", None
)

2.3 DataModule：数据管理最佳实践

LightningDataModule 将数据集、数据变换和数据加载器封装在一起，实现数据逻辑的重用。

基础 DataModule

from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, random_split

class MNISTDataModule(LightningDataModule):
    def __init__(
        self,
        data_dir: str = "./",
        batch_size: int = 32,
        num_workers: int = 4,
    ):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])

    def prepare_data(self):
        """下载数据（只在主进程执行一次）"""
        datasets.MNIST(self.data_dir, train=True, download=True)
        datasets.MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage: str):
        """设置数据集（每个进程执行）"""
        if stage == "fit" or stage is None:
            mnist_full = datasets.MNIST(
                self.data_dir, train=True, transform=self.transform
            )
            self.mnist_train, self.mnist_val = random_split(
                mnist_full, [55000, 5000]
            )

        if stage == "test" or stage is None:
            self.mnist_test = datasets.MNIST(
                self.data_dir, train=False, transform=self.transform
            )

        if stage == "predict" or stage is None:
            self.mnist_predict = datasets.MNIST(
                self.data_dir, train=False, transform=self.transform
            )

    def train_dataloader(self):
        return DataLoader(
            self.mnist_train,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=True,
            pin_memory=True,  # 加速 GPU 训练
        )

    def val_dataloader(self):
        return DataLoader(
            self.mnist_val,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
        )

    def test_dataloader(self):
        return DataLoader(
            self.mnist_test,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
        )

    def predict_dataloader(self):
        return DataLoader(
            self.mnist_predict,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
        )


# 使用 DataModule
if __name__ == "__main__":
    datamodule = MNISTDataModule(batch_size=64)
    model = MNISTModel()

    trainer = Trainer(max_epochs=5)
    trainer.fit(model, datamodule=datamodule)

高级 DataModule：自定义数据集

import pandas as pd
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, csv_file, transform=None):
        self.data_frame = pd.read_csv(csv_file)
        self.transform = transform

    def __len__(self):
        return len(self.data_frame)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        # 加载数据
        img_name = self.data_frame.iloc[idx, 0]
        image = io.imread(img_name)
        label = self.data_frame.iloc[idx, 1]

        if self.transform:
            image = self.transform(image)

        return image, label


class CustomDataModule(LightningDataModule):
    def __init__(
        self,
        train_csv: str,
        val_csv: str,
        test_csv: str,
        batch_size: int = 32,
        num_workers: int = 4,
    ):
        super().__init__()
        self.train_csv = train_csv
        self.val_csv = val_csv
        self.test_csv = test_csv
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage: str):
        if stage == "fit":
            self.train_dataset = CustomDataset(
                self.train_csv,
                transform=transforms.Compose([
                    transforms.ToTensor(),
                    transforms.RandomHorizontalFlip(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                ])
            )
            self.val_dataset = CustomDataset(
                self.val_csv,
                transform=transforms.Compose([
                    transforms.ToTensor(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                ])
            )

        if stage == "test":
            self.test_dataset = CustomDataset(
                self.test_csv,
                transform=transforms.Compose([
                    transforms.ToTensor(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                ])
            )

    def train_dataloader(self):
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=True,
        )

    def val_dataloader(self):
        return DataLoader(
            self.val_dataset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
        )

    def test_dataloader(self):
        return DataLoader(
            self.test_dataset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
        )

2.4 Callbacks：自定义训练流程

Callbacks 允许你在训练的关键点插入自定义逻辑，而无需修改 LightningModule。

内置 Callbacks

1. ModelCheckpoint - 模型检查点

from pytorch_lightning.callbacks import ModelCheckpoint

checkpoint_callback = ModelCheckpoint(
    dirpath="checkpoints/",
    filename="model-{epoch:02d}-{val_loss:.2f}",
    monitor="val_loss",        # 监控的指标
    save_top_k=3,              # 保存最好的 3 个模型
    mode="min",                # "min" (越小越好) or "max" (越大越好)
    save_last=True,            # 保存最后一个 epoch 的模型
    save_weights_only=False,   # False 保存整个模型，True 只保存权重
    every_n_train_steps=50,    # 每 N 步保存一次
)

trainer = Trainer(callbacks=[checkpoint_callback])

2. EarlyStopping - 早期停止

from pytorch_lightning.callbacks import EarlyStopping

early_stop_callback = EarlyStopping(
    monitor="val_loss",
    patience=5,               # 容忍 5 个 epoch 不改进
    verbose=True,
    mode="min",
    min_delta=0.001,          # 改进小于此值不算改进
    check_on_train_epoch_end=True,
)

trainer = Trainer(callbacks=[early_stop_callback])

3. LearningRateMonitor - 学习率监控

from pytorch_lightning.callbacks import LearningRateMonitor

lr_monitor = LearningRateMonitor(
    logging_interval="step",  # "step" or "epoch"
)

trainer = Trainer(callbacks=[lr_monitor])

4. ModelSummary - 模型摘要

from pytorch_lightning.callbacks import ModelSummary

model_summary = ModelSummary(
    max_depth=3,  # 显示的层数深度
)

trainer = Trainer(callbacks=[model_summary])

自定义 Callback

from pytorch_lightning.callbacks import Callback

class MyCustomCallback(Callback):
    def on_train_start(self, trainer, pl_module):
        print("Training is starting!")

    def on_train_end(self, trainer, pl_module):
        print("Training is ending!")

    def on_validation_end(self, trainer, pl_module):
        # 每个 epoch 验证结束后打印一些指标
        print(f"Epoch {trainer.current_epoch} validation ended")
        print(f"Current val_loss: {trainer.callback_metrics['val_loss']}")

    def on_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        # 每个 batch 结束后执行
        if batch_idx % 100 == 0:
            print(f"Batch {batch_idx} ended")

    def on_exception(self, trainer, pl_module, exception):
        # 训练异常时执行
        print(f"Exception occurred: {exception}")


# 组合多个 Callbacks
trainer = Trainer(
    callbacks=[
        ModelCheckpoint(monitor="val_loss", save_top_k=1),
        EarlyStopping(monitor="val_loss", patience=5),
        LearningRateMonitor(logging_interval="epoch"),
        MyCustomCallback(),
    ]
)

高级 Callback：梯度裁剪和累积

from pytorch_lightning.callbacks import GradientAccumulationScheduler

# 不同 epoch 使用不同的累积步数
accumulation_scheduler = GradientAccumulationScheduler(
    scheduling={0: 4, 5: 8, 10: 16}  # epoch 0-4: 4, 5-9: 8, 10+: 16
)

trainer = Trainer(callbacks=[accumulation_scheduler])

第三部分：高级特性

3.1 多 GPU 训练策略

Lightning 支持多种并行策略，从单机多卡到多节点分布式训练。

策略对比

策略	全称	适用场景	通信开销	代码改动
dp	DataParallel	单机多卡（小模型）	高	无
ddp	DistributedDataParallel	单机多卡/多节点（推荐）	中	无
ddp_spawn	DDP with spawn	Jupyter/Interactive	中	无
fsdp	FullyShardedDataParallel	大模型（分片）	低	无
deepspeed	DeepSpeed	超大模型	低	需配置

DataParallel (DP)

# 最简单，但性能较差（不推荐用于生产）
trainer = Trainer(
    accelerator="gpu",
    devices=4,
    strategy="dp",  # DataParallel
)

特点：

✅ 最简单，无需修改代码
❌ 单进程多线程，GPU 利用率不均衡
❌ 显存占用高（每个 GPU 存储完整模型）
❌ 通信开销大

DistributedDataParallel (DDP) - 推荐

# 最佳性能（推荐）
trainer = Trainer(
    accelerator="gpu",
    devices=4,
    strategy="ddp",  # DistributedDataParallel
    num_nodes=1,     # 单节点
)

# 多节点分布式训练
trainer = Trainer(
    accelerator="gpu",
    devices=8,       # 每个节点 8 张卡
    num_nodes=4,     # 4 个节点
    strategy="ddp",
)

启动多节点训练：

# 节点 0 (主节点)
python train.py --accelerator gpu --devices 8 --num_nodes 4 --node_rank 0

# 节点 1
python train.py --accelerator gpu --devices 8 --num_nodes 4 --node_rank 1

# 节点 2
python train.py --accelerator gpu --devices 8 --num_nodes 4 --node_rank 2

# 节点 3
python train.py --accelerator gpu --devices 8 --num_nodes 4 --node_rank 3

特点：

✅ 多进程并行，GPU 利用率高
✅ 显存占用低（每个 GPU 只存储模型切片）
✅ 通信效率高
✅ 可扩展到多节点

FullyShardedDataParallel (FSDP) - 大模型优化

# 用于超大模型（无法放入单个 GPU）
trainer = Trainer(
    accelerator="gpu",
    devices=8,
    strategy="fsdp",  # FullyShardedDataParallel
    precision="bf16", # 推荐 BF16
)

FSDP 高级配置：

from pytorch_lightning.strategies import FSDPStrategy

# 自定义 FSDP 配置
strategy = FSDPStrategy(
    sharding_strategy="FULL_SHARD",  # "FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD"
    cpu_offload=False,               # CPU 卸载
    auto_wrap_policy=...             # 自动包装策略
)

trainer = Trainer(
    accelerator="gpu",
    devices=8,
    strategy=strategy,
)

DeepSpeed 集成

from pytorch_lightning.strategies import DeepSpeedStrategy

# DeepSpeed 配置
strategy = DeepSpeedStrategy(
    stage=2,  # ZeRO Stage 0/1/2/3
    offload_optimizer=True,  # CPU 卸载优化器状态
    offload_parameters=False,
)

trainer = Trainer(
    accelerator="gpu",
    devices=8,
    strategy=strategy,
)

ZeRO Stage 说明：

Stage 0：无分片（等同于 DDP）
Stage 1：优化器状态分片
Stage 2：优化器状态 + 梯度分片
Stage 3：优化器状态 + 梯度 + 参数分片

3.2 混合精度训练 (AMP)

混合精度训练使用 16 位浮点数（FP16）或 BFloat16（BF16）进行计算，显著提升训练速度并减少显存占用。

FP16 混合精度

# FP16 混合精度
trainer = Trainer(
    precision=16,
    accelerator="gpu",
    devices=1,
)

# 完整配置
trainer = Trainer(
    precision=16,
    amp_backend="native",  # "native" (PyTorch AMP) or "apex"
    amp_level="O1",        # Apex O1/O2 (仅当使用 Apex 时)
)

BF16（推荐，2025+）

# BFloat16 - 更稳定的数值范围（推荐）
trainer = Trainer(
    precision="bf16",
    accelerator="gpu",
    devices=1,
)

FP16 vs BF16 对比：

特性	FP16	BF16
数值范围	有限（需要 Loss Scaling）	与 FP32 相同
稳定性	需要调参	更稳定
硬件要求	Volta+	Ampere+ (A100, RTX 30xx)
推荐场景	老款 GPU	新款 GPU (A100, 3090, 4090)

自定义混合精度配置

model = LitModel()
trainer = Trainer(
    precision=16,
    amp_backend="native",
    amp_type="bfloat16",  # "float16" or "bfloat16"
)
trainer.fit(model)

3.3 梯度操作

梯度累积

梯度累积允许在多次小 batch 计算后再更新参数，模拟大 batch size 训练：

# 基础用法
trainer = Trainer(
    accumulate_grad_batches=4,  # 每 4 个 batch 更新一次参数
)

# 动态累积步数（不同 epoch 使用不同步数）
trainer = Trainer(
    accumulate_grad_batches={
        0: 4,   # 前 5 个 epoch: 累积 4 步
        5: 8,   # 第 5-9 个 epoch: 累积 8 步
        10: 16, # 第 10+ 个 epoch: 累积 16 步
    }
)

有效 Batch Size 计算：

actual_batch_size = batch_size * accumulate_grad_batches * num_gpus * num_nodes
# 例如: 32 * 4 * 4 * 1 = 512

梯度裁剪

# 按范数裁剪
trainer = Trainer(
    gradient_clip_val=1.0,           # 裁剪阈值
    gradient_clip_algorithm="norm",  # "norm" or "value"
)

# 按值裁剪
trainer = Trainer(
    gradient_clip_val=0.5,
    gradient_clip_algorithm="value",  # 每个梯度元素裁剪到 [-0.5, 0.5]
)

在 LightningModule 中自定义：

class LitModel(pl.LightningModule):
    def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_val, gradient_clip_algorithm):
        # 自定义梯度裁剪逻辑
        self.clip_gradients(
            optimizer,
            gradient_clip_val=gradient_clip_val,
            gradient_clip_algorithm=gradient_clip_algorithm,
        )

3.4 超参数调优

Lightning 提供了内置的超参数调优工具。

学习率查找器

# 自动查找最佳学习率
trainer = Trainer(
    auto_lr_find=True,  # 启用自动学习率查找
)
trainer.tune(model)  # 在训练前调用

# 可视化学习率查找结果
model = LitModel()
trainer = Trainer(auto_lr_find=True)
lr_finder = trainer.tuner.lr_find(model)
fig = lr_finder.plot(suggest=True)
fig.show()

# 获取建议的学习率
new_lr = lr_finder.suggest()
model.hparams.learning_rate = new_lr

批处理大小查找器

trainer = Trainer(
    auto_scale_batch_size="binsearch",  # "binsearch" or "power"
)
trainer.tune(model)  # 自动查找最大可用 batch size

完整自动调优

# 综合调优
model = LitModel(learning_rate=0.001)  # 初始学习率
trainer = Trainer(
    auto_lr_find=True,
    auto_scale_batch_size="binsearch",
    max_epochs=50,
)

# 1. 查找最佳学习率和批大小
trainer.tune(model)

# 2. 使用优化后的超参数训练
trainer.fit(model)

3.5 性能优化技巧

数据加载优化

# DataModule 优化配置
class OptimizedDataModule(LightningDataModule):
    def train_dataloader(self):
        return DataLoader(
            self.train_dataset,
            batch_size=64,
            num_workers=4,        # 数据加载进程数
            pin_memory=True,      # 加速 CPU->GPU 传输
            persistent_workers=True,  # 保持 worker 进程
            prefetch_factor=2,    # 每个 worker 预取 2 个 batch
        )

模型编译（PyTorch 2.0+）

class LitModel(pl.LightningModule):
    def configure_model(self):
        """PyTorch 2.0+ 编译优化"""
        if torch.__version__ >= "2.0":
            self.model = torch.compile(self.model, mode="max-autotune")

其他优化选项

trainer = Trainer(
    # 1. 禁用进度条（加速训练）
    enable_progress_bar=False,

    # 2. 减少日志频率
    log_every_n_steps=100,

    # 3. 使用 Deterministic 训练（可复现性）
    deterministic=True,

    # 4. Benchmark 模式（加速，但不可复现）
    benchmark=True,

    # 5. 性能分析器
    profiler="advanced",  # "simple", "advanced", or PyTorch Profiler
)

第四部分：实战案例

4.1 计算机视觉：图像分类（ResNet）

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

class ResNetClassifier(pl.LightningModule):
    def __init__(
        self,
        num_classes: int = 10,
        learning_rate: float = 1e-3,
        pretrained: bool = True,
    ):
        super().__init__()
        self.save_hyperparameters()

        # 加载预训练 ResNet
        self.model = models.resnet50(pretrained=pretrained)

        # 替换最后一层（适配分类数）
        num_features = self.model.fc.in_features
        self.model.fc = nn.Linear(num_features, num_classes)

        # 损失函数
        self.criterion = nn.CrossEntropyLoss()

        # 评估指标
        self.train_accuracy = []
        self.val_accuracy = []
        self.test_accuracy = []

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)

        # 计算准确率
        preds = torch.argmax(y_hat, dim=1)
        acc = (preds == y).float().mean()

        # 记录指标
        self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=True)
        self.log("train_acc", acc, prog_bar=True, on_step=True, on_epoch=True)

        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)

        preds = torch.argmax(y_hat, dim=1)
        acc = (preds == y).float().mean()

        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)

        return {"val_loss": loss, "val_acc": acc}

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)

        preds = torch.argmax(y_hat, dim=1)
        acc = (preds == y).float().mean()

        self.log("test_loss", loss)
        self.log("test_acc", acc)

        return {"test_loss": loss, "test_acc": acc}

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=self.hparams.learning_rate,
            weight_decay=0.01,
        )

        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer,
            T_0=10,
            T_mult=2,
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "epoch",
                "frequency": 1,
            },
        }

    def on_train_epoch_end(self):
        # 每个 epoch 结束后记录学习率
        opt = self.optimizers()
        current_lr = opt.param_groups[0]["lr"]
        self.log("lr", current_lr)


# CIFAR-10 DataModule
class CIFAR10DataModule(pl.LightningDataModule):
    def __init__(self, data_dir="./", batch_size=32, num_workers=4):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers

        # 训练数据增强
        self.train_transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])

        # 测试数据标准化
        self.test_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])

    def prepare_data(self):
        datasets.CIFAR10(self.data_dir, train=True, download=True)
        datasets.CIFAR10(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        if stage == "fit" or stage is None:
            cifar_full = datasets.CIFAR10(
                self.data_dir, train=True, transform=self.train_transform
            )
            self.cifar_train, self.cifar_val = torch.utils.data.random_split(
                cifar_full, [45000, 5000]
            )

        if stage == "test" or stage is None:
            self.cifar_test = datasets.CIFAR10(
                self.data_dir, train=False, transform=self.test_transform
            )

    def train_dataloader(self):
        return DataLoader(
            self.cifar_train,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=True,
            pin_memory=True,
        )

    def val_dataloader(self):
        return DataLoader(
            self.cifar_val,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
        )

    def test_dataloader(self):
        return DataLoader(
            self.cifar_test,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
        )


# 训练脚本
if __name__ == "__main__":
    from pytorch_lightning.loggers import TensorBoardLogger
    from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping

    # 初始化
    model = ResNetClassifier(num_classes=10, learning_rate=1e-3)
    datamodule = CIFAR10DataModule(batch_size=64, num_workers=4)

    # Callbacks
    checkpoint_callback = ModelCheckpoint(
        dirpath="checkpoints/resnet_cifar10/",
        filename="resnet-{epoch:02d}-{val_acc:.4f}",
        monitor="val_acc",
        mode="max",
        save_top_k=3,
        save_last=True,
    )

    early_stop_callback = EarlyStopping(
        monitor="val_acc",
        patience=10,
        mode="max",
    )

    # Logger
    logger = TensorBoardLogger("logs/", name="resnet_cifar10")

    # Trainer
    trainer = pl.Trainer(
        max_epochs=100,
        accelerator="auto",
        devices="auto",
        precision=16,  # 混合精度
        callbacks=[checkpoint_callback, early_stop_callback],
        logger=logger,
        gradient_clip_val=1.0,
        accumulate_grad_batches=2,
    )

    # 训练
    trainer.fit(model, datamodule=datamodule)

    # 测试
    trainer.test(model, datamodule=datamodule)

4.2 自然语言处理：文本分类（BERT）

import pytorch_lightning as pl
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from transformers import BertModel, BertTokenizer, AdamW

# 自定义数据集
class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            return_token_type_ids=False,
            padding="max_length",
            truncation=True,
            return_attention_mask=True,
            return_tensors="pt",
        )

        return {
            "input_ids": encoding["input_ids"].flatten(),
            "attention_mask": encoding["attention_mask"].flatten(),
            "labels": torch.tensor(label, dtype=torch.long),
        }


# BERT 分类器
class BertClassifier(pl.LightningModule):
    def __init__(
        self,
        num_classes: int = 2,
        learning_rate: float = 2e-5,
        model_name: str = "bert-base-uncased",
    ):
        super().__init__()
        self.save_hyperparameters()

        # BERT 模型
        self.bert = BertModel.from_pretrained(model_name)

        # 分类头
        self.classifier = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(self.bert.config.hidden_size, num_classes),
        )

        # 损失函数
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        logits = self.classifier(pooled_output)
        return logits

    def training_step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]

        logits = self(input_ids, attention_mask)
        loss = self.criterion(logits, labels)

        preds = torch.argmax(logits, dim=1)
        acc = (preds == labels).float().mean()

        self.log("train_loss", loss, prog_bar=True)
        self.log("train_acc", acc, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]

        logits = self(input_ids, attention_mask)
        loss = self.criterion(logits, labels)

        preds = torch.argmax(logits, dim=1)
        acc = (preds == labels).float().mean()

        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)

    def test_step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]

        logits = self(input_ids, attention_mask)
        loss = self.criterion(logits, labels)

        preds = torch.argmax(logits, dim=1)
        acc = (preds == labels).float().mean()

        self.log("test_loss", loss)
        self.log("test_acc", acc)

    def configure_optimizers(self):
        return AdamW(self.parameters(), lr=self.hparams.learning_rate)


# DataModule
class TextDataModule(pl.LightningDataModule):
    def __init__(
        self,
        train_texts,
        train_labels,
        val_texts,
        val_labels,
        test_texts=None,
        test_labels=None,
        model_name="bert-base-uncased",
        batch_size=32,
        num_workers=4,
    ):
        super().__init__()
        self.train_texts = train_texts
        self.train_labels = train_labels
        self.val_texts = val_texts
        self.val_labels = val_labels
        self.test_texts = test_texts
        self.test_labels = test_labels
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage=None):
        if stage == "fit" or stage is None:
            self.train_dataset = TextClassificationDataset(
                self.train_texts, self.train_labels, self.tokenizer
            )
            self.val_dataset = TextClassificationDataset(
                self.val_texts, self.val_labels, self.tokenizer
            )

        if stage == "test" or stage is None:
            if self.test_texts is not None:
                self.test_dataset = TextClassificationDataset(
                    self.test_texts, self.test_labels, self.tokenizer
                )

    def train_dataloader(self):
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=True,
        )

    def val_dataloader(self):
        return DataLoader(
            self.val_dataset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
        )

    def test_dataloader(self):
        if hasattr(self, "test_dataset"):
            return DataLoader(
                self.test_dataset,
                batch_size=self.batch_size,
                num_workers=self.num_workers,
            )
        return None


# 使用示例
if __name__ == "__main__":
    # 模拟数据
    train_texts = ["This is a positive sentence."]*100 + ["This is a negative sentence."]*100
    train_labels = [1]*100 + [0]*100
    val_texts = ["Another positive example."]*20 + ["Another negative example."]*20
    val_labels = [1]*20 + [0]*20

    # 初始化
    model = BertClassifier(num_classes=2)
    datamodule = TextDataModule(
        train_texts, train_labels,
        val_texts, val_labels,
        batch_size=16,
    )

    # 训练
    trainer = pl.Trainer(
        max_epochs=3,
        accelerator="auto",
        devices="auto",
        precision=16,
    )
    trainer.fit(model, datamodule=datamodule)

4.3 推荐系统：神经协同过滤

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F

class NeuralCollaborativeFiltering(pl.LightningModule):
    def __init__(
        self,
        num_users: int,
        num_items: int,
        embedding_dim: int = 64,
        hidden_dims: list = [128, 64],
        learning_rate: float = 1e-3,
        dropout: float = 0.2,
    ):
        super().__init__()
        self.save_hyperparameters()

        # 嵌入层
        self.user_embedding = nn.Embedding(num_users, embedding_dim)
        self.item_embedding = nn.Embedding(num_items, embedding_dim)

        # MLP 层
        input_dim = embedding_dim * 2
        layers = []
        prev_dim = input_dim
        for dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))
            prev_dim = dim
        layers.append(nn.Linear(prev_dim, 1))
        self.mlp = nn.Sequential(*layers)

        # 初始化权重
        self._init_weights()

    def _init_weights(self):
        nn.init.normal_(self.user_embedding.weight, std=0.01)
        nn.init.normal_(self.item_embedding.weight, std=0.01)

    def forward(self, user_ids, item_ids):
        user_embed = self.user_embedding(user_ids)  # [batch, embedding_dim]
        item_embed = self.item_embedding(item_ids)  # [batch, embedding_dim]

        # 拼接嵌入向量
        concat = torch.cat([user_embed, item_embed], dim=1)  # [batch, 2*embedding_dim]

        # MLP 预测
        output = self.mlp(concat)  # [batch, 1]
        return output.squeeze()

    def training_step(self, batch, batch_idx):
        user_ids, item_ids, ratings = batch
        predictions = self(user_ids, item_ids)

        loss = F.mse_loss(predictions, ratings.float())

        self.log("train_loss", loss, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        user_ids, item_ids, ratings = batch
        predictions = self(user_ids, item_ids)

        loss = F.mse_loss(predictions, ratings.float())

        # RMSE
        rmse = torch.sqrt(loss)

        self.log("val_loss", loss, prog_bar=True)
        self.log("val_rmse", rmse, prog_bar=True)

    def test_step(self, batch, batch_idx):
        user_ids, item_ids, ratings = batch
        predictions = self(user_ids, item_ids)

        loss = F.mse_loss(predictions, ratings.float())
        rmse = torch.sqrt(loss)

        self.log("test_loss", loss)
        self.log("test_rmse", rmse)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(
            self.parameters(),
            lr=self.hparams.learning_rate,
        )

        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            mode="min",
            factor=0.5,
            patience=5,
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val_loss",
            },
        }


# 数据集示例
class MovieLensDataset(torch.utils.data.Dataset):
    def __init__(self, user_ids, item_ids, ratings):
        self.user_ids = torch.LongTensor(user_ids)
        self.item_ids = torch.LongTensor(item_ids)
        self.ratings = torch.FloatTensor(ratings)

    def __len__(self):
        return len(self.ratings)

    def __getitem__(self, idx):
        return (
            self.user_ids[idx],
            self.item_ids[idx],
            self.ratings[idx],
        )


# 训练示例
if __name__ == "__main__":
    # 模拟数据
    num_users = 1000
    num_items = 500
    num_samples = 10000

    user_ids = torch.randint(0, num_users, (num_samples,))
    item_ids = torch.randint(0, num_items, (num_samples,))
    ratings = torch.rand(num_samples) * 4 + 1  # 1-5 评分

    # 划分数据集
    split = int(0.8 * num_samples)
    train_dataset = MovieLensDataset(
        user_ids[:split].numpy(),
        item_ids[:split].numpy(),
        ratings[:split].numpy(),
    )
    val_dataset = MovieLensDataset(
        user_ids[split:].numpy(),
        item_ids[split:].numpy(),
        ratings[split:].numpy(),
    )

    train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=256)

    # 初始化模型
    model = NeuralCollaborativeFiltering(
        num_users=num_users,
        num_items=num_items,
        embedding_dim=64,
        hidden_dims=[128, 64],
    )

    # 训练
    trainer = pl.Trainer(
        max_epochs=50,
        accelerator="auto",
        devices="auto",
        precision=16,
    )
    trainer.fit(model, train_loader, val_loader)

第五部分：最佳实践与优化

5.1 代码组织最佳实践

项目结构推荐

project/
├── data/
│   ├── processed/
│   └── raw/
├── models/
│   ├── __init__.py
│   ├── lit_model.py
│   └── datamodule.py
├── logs/
├── checkpoints/
├── configs/
│   └── config.yaml
├── src/
│   ├── __init__.py
│   ├── train.py
│   └── utils.py
└── tests/

配置管理（Hydra + Lightning）

# configs/config.yaml
model:
  learning_rate: 0.001
  hidden_dim: 128
  dropout: 0.2

data:
  batch_size: 64
  num_workers: 4

trainer:
  max_epochs: 100
  accelerator: auto
  devices: auto
  precision: 16

# train.py
import pytorch_lightning as pl
import hydra
from omegaconf import DictConfig

@hydra.main(config_path="configs", config_name="config")
def train(cfg: DictConfig):
    model = LitModel(**cfg.model)
    datamodule = MyDataModule(**cfg.data)

    trainer = pl.Trainer(**cfg.trainer)
    trainer.fit(model, datamodule=datamodule)

if __name__ == "__main__":
    train()

5.2 性能优化清单

数据加载优化

# ✅ 好的实践
dataloader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,              # 多进程加载
    pin_memory=True,            # 加速 CPU->GPU 传输
    persistent_workers=True,    # 保持 worker 进程
    prefetch_factor=2,          # 预取 batch
)

# ❌ 避免的做法
dataloader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=0,              # 单进程加载（慢）
    pin_memory=False,
)

训练优化

trainer = Trainer(
    # 1. 混合精度
    precision=16,

    # 2. 梯度累积
    accumulate_grad_batches=4,

    # 3. 梯度裁剪
    gradient_clip_val=1.0,
    gradient_clip_algorithm="norm",

    # 4. 关闭不必要的功能（加速训练）
    enable_progress_bar=False,    # 生产环境可关闭
    enable_model_summary=False,
    log_every_n_steps=100,        # 减少日志频率

    # 5. 性能分析（调试时启用）
    profiler="simple",  # 或 "advanced"
)

内存优化

# ✅ 使用梯度检查点（节省显存）
model = LitModel()
model = torch.utils.checkpoint.checkpoint_sequential(model, segments=4)

# ✅ 清理不必要的缓存
class LitModel(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        loss = self(batch)
        # 手动清理中间变量
        del intermediate_results
        torch.cuda.empty_cache()  # 谨慎使用（性能开销）
        return loss

5.3 常见陷阱和解决方案

陷阱 1：忘记 `self.log()` 中的 `on_epoch=True`

# ❌ 错误
def training_step(self, batch, batch_idx):
    loss = self(batch)
    self.log("train_loss", loss)  # 只记录每步的值
    return loss

# ✅ 正确
def training_step(self, batch, batch_idx):
    loss = self(batch)
    self.log("train_loss", loss, on_step=True, on_epoch=True)  # 同时记录步和 epoch
    return loss

陷阱 2：DataModule 中 `prepare_data` 使用不当

# ❌ 错误：在 setup 中下载数据（会被每个进程调用）
def setup(self, stage=None):
    self.data = datasets.MNIST(".", download=True)  # 每个进程都下载！

# ✅ 正确：在 prepare_data 中下载（只调用一次）
def prepare_data(self):
    datasets.MNIST(".", download=True)  # 只在主进程下载

def setup(self, stage=None):
    self.data = datasets.MNIST(".", download=False)

陷阱 3：分布式训练时的数据采样

# ❌ 错误：使用 DistributedSampler 但设置 shuffle=True
train_loader = DataLoader(
    dataset,
    shuffle=True,  # 与 DistributedSampler 冲突
)

# ✅ 正确：DistributedSampler 会自动处理 shuffle
train_loader = DataLoader(
    dataset,
    sampler=DistributedSampler(dataset, shuffle=True),
)

陷阱 4：验证集数据增强

# ❌ 错误：验证集使用数据增强
class MyDataModule(LightningDataModule):
    def setup(self, stage=None):
        if stage == "fit":
            self.train_dataset = Dataset(transform=augmentation_transform)
            self.val_dataset = Dataset(transform=augmentation_transform)  # 错误！

# ✅ 正确：验证集只做标准化
class MyDataModule(LightningDataModule):
    def setup(self, stage=None):
        if stage == "fit":
            self.train_dataset = Dataset(transform=augmentation_transform)
            self.val_dataset = Dataset(transform=normalize_transform)

5.4 生态系统集成

Weights & Biases (W&B)

import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger

# W&B Logger
wandb_logger = WandbLogger(
    project="my-project",
    entity="my-team",
    name="experiment-1",
    save_dir="logs/",
    offline=False,  # True 用于离线模式
)

# 记录额外信息
wandb_logger.experiment.config["batch_size"] = 32
wandb_logger.experiment.config["learning_rate"] = 0.001

trainer = pl.Trainer(
    logger=wandb_logger,
    max_epochs=100,
)
trainer.fit(model)

# 记录模型
wandb_logger.experiment.watch(model, log="all", log_freq=100)

TensorBoard

from pytorch_lightning.loggers import TensorBoardLogger

logger = TensorBoardLogger(
    save_dir="logs/",
    name="my_model",
    version="version_1",  # None 自动递增
    default_hp_metric=False,
)

trainer = pl.Trainer(logger=logger)
trainer.fit(model)

# 启动 TensorBoard
# tensorboard --logdir logs/

MLflow

from pytorch_lightning.loggers import MLFlowLogger

mlflow_logger = MLFlowLogger(
    experiment_name="my-experiment",
    tracking_uri="file:///./mlruns",
    tags={"team": "research", "project": "nlp"},
)

trainer = pl.Trainer(logger=mlflow_logger)
trainer.fit(model)

Neptune AI

import neptune
from pytorch_lightning.loggers import NeptuneLogger

neptune_logger = NeptuneLogger(
    api_key="YOUR_API_KEY",
    project_name="workspace/project",
    tags=["pytorch-lightning", "experiment"],
)

trainer = pl.Trainer(logger=neptune_logger)
trainer.fit(model)

组合多个 Logger

from pytorch_lightning.loggers import WandbLogger, TensorBoardLogger

# 同时使用多个 logger
wandb_logger = WandbLogger(project="my-project")
tensorboard_logger = TensorBoardLogger(save_dir="logs/")

trainer = pl.Trainer(
    logger=[wandb_logger, tensorboard_logger],
)

第六部分：生产部署

6.1 模型导出

导出到 TorchScript

class LitModel(pl.LightningModule):
    def forward(self, x):
        return self.model(x)

# 训练后导出
model = LitModel.load_from_checkpoint("checkpoints/model.ckpt")
model.eval()

# 方法 1：Tracing
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_traced.pt")

# 方法 2：Scripting（更灵活）
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

# 加载和推理
loaded_model = torch.jit.load("model_scripted.pt")
output = loaded_model(example_input)

导出到 ONNX

# 导出为 ONNX
model = LitModel.load_from_checkpoint("checkpoints/model.ckpt")
model.eval()

example_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    example_input,
    "model.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"},
    },
)

6.2 TorchServe 部署

# 1. 安装 TorchServe
# pip install torchserve torch-model-archiver torch-workflow-archiver

# 2. 创建模型处理器（handler.py）
class ModelHandler(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = LitModel.load_from_checkpoint("model.ckpt")
        self.model.eval()

    def forward(self, x):
        with torch.no_grad():
            return self.model(x)

# 3. 导出模型
model = ModelHandler()
scripted_model = torch.jit.script(model)
scripted_model.save("model_handler.pt")

# 4. 打包模型
# torch-model-archiver --model-name my_model \
#   --version 1.0 \
#   --serialized-file model_handler.pt \
#   --handler handler.py \
#   --extra-files config.json

# 5. 启动 TorchServe
# torchserve --start --ncs --model-store model_store --models my_model=my_model.mar

6.3 Lightning Apps（2025 新特性）

Lightning Apps 是 Lightning AI 的云原生应用框架，简化模型部署。

import lightning.app as L
import lightning as L
from lightning.app.components import ServeTrainer, PythonServer
from lightning.app.frontend import StreamlitFrontend

# 定义训练组件
class MyTrainingComponent(L.LightningWork):
    def run(self):
        model = LitModel()
        trainer = L.Trainer(max_epochs=10)
        trainer.fit(model)
        return "Training completed!"

# 定义推理服务
class MyModelServer(PythonServer):
    def __init__(self, model_path):
        super().__init__()
        self.model_path = model_path

    def setup(self):
        self.model = LitModel.load_from_checkpoint(self.model_path)
        self.model.eval()

    def predict(self, data):
        with torch.no_grad():
            return self.model(data)

# Lightning App
class MyApp(L.LightningFlow):
    def __init__(self):
        super().__init__()
        self.training = MyTrainingComponent()
        self.serving = MyModelServer("model.ckpt")

    def run(self):
        self.training.run()
        if self.training.has_succeeded:
            self.serving.run()

# Frontend（可选）
import streamlit as st

def render(self):
    st.title("My Model")
    user_input = st.text_input("Enter text")
    if st.button("Predict"):
        result = self.serving.predict(user_input)
        st.write(result)

# 运行 App
app = L.LightningApp(MyApp())

6.4 生产环境配置

检查点管理

from pytorch_lightning.callbacks import ModelCheckpoint

# 生产级检查点配置
checkpoint_callback = ModelCheckpoint(
    dirpath="production/checkpoints/",
    filename="model-{epoch:02d}-{val_loss:.4f}",
    monitor="val_loss",
    mode="min",
    save_top_k=1,        # 只保存最佳模型
    save_last=True,      # 保存最后一个 checkpoint
    save_weights_only=True,  # 只保存权重（节省空间）
    every_n_train_steps=500,
)

# 自动删除旧检查点
trainer = Trainer(
    callbacks=[
        checkpoint_callback,
    ],
    default_root_dir="production/",
)

模型版本管理

import datetime

class VersionedModelCheckpoint(ModelCheckpoint):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def _save_checkpoint(self, trainer, filepath):
        # 添加版本信息
        version = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        filepath = filepath.replace(".ckpt", f"_v{version}.ckpt")
        super()._save_checkpoint(trainer, filepath)

# 使用
checkpoint_callback = VersionedModelCheckpoint(
    dirpath="checkpoints/",
    filename="model",
    monitor="val_loss",
)

第七部分：总结与展望

7.1 Lightning vs 其他框架

框架	灵活性	易用性	分布式训练	生态系统	学习曲线
PyTorch Lightning	⭐⭐⭐⭐	⭐⭐⭐⭐⭐	⭐⭐⭐⭐⭐	⭐⭐⭐⭐	低
Vanilla PyTorch	⭐⭐⭐⭐⭐	⭐⭐	⭐⭐	⭐⭐⭐⭐⭐	高
Hugging Face Trainer	⭐⭐⭐	⭐⭐⭐⭐⭐	⭐⭐⭐⭐	⭐⭐⭐⭐⭐	低
Keras	⭐⭐	⭐⭐⭐⭐⭐	⭐⭐⭐	⭐⭐⭐⭐⭐	极低
FastAI	⭐⭐⭐	⭐⭐⭐⭐	⭐⭐	⭐⭐⭐	低

7.2 何时使用 Lightning？

✅ 推荐使用 Lightning 的场景：

需要多 GPU/多节点训练
研究项目需要快速迭代
团队协作，需要标准化代码结构
需要完整的实验跟踪和日志记录
迁移 vanilla PyTorch 项目到生产环境
教育和培训（代码更清晰）

❌ 可能不需要 Lightning 的场景：

极其简单的模型（单文件脚本）
需要对训练循环进行极度的自定义控制
已有成熟的训练流程，迁移成本高

7.3 2025 及未来展望

2025 年路线图

更好的分布式训练
- 完整的 FSDP 支持
- 多模态大模型训练优化
深度学习编译器集成
- 与 torch.compile 无缝集成
- 自动优化训练流程
云原生部署
- Lightning Apps 的持续改进
- Serverless 训练和推理
更好的可观测性
- 内置性能分析工具
- 实时监控和调试
生态系统扩展
- 更多官方 Callbacks
- 与更多工具集成

7.4 学习资源

官方资源

推荐教程

入门
- Lightning in 15 Minutes
- 官方教程集合
进阶
- 高级训练技巧
- 性能优化指南
实战
- Lightning Examples
- Lightning Bolts（预训练模型）

社区资源

附录：快速参考

A. 常用命令速查

# 安装
pip install pytorch-lightning

# 查看版本
python -c "import pytorch_lightning as pl; print(pl.__version__)"

# 启动 TensorBoard
tensorboard --logdir logs/

# 多 GPU 训练
python train.py --accelerator gpu --devices 4 --strategy ddp

# 混合精度训练
python train.py --precision 16

# 恢复训练
python train.py --ckpt_path checkpoints/last.ckpt

B. 代码模板速查

# 最小 Lightning 模板
import pytorch_lightning as pl
import torch.nn as nn

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 1)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.functional.mse_loss(y_hat, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)

# 训练
model = LitModel()
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, train_dataloader, val_dataloader)

C. 常见配置速查

# 多 GPU + 混合精度
trainer = Trainer(
    accelerator="gpu",
    devices=4,
    strategy="ddp",
    precision=16,
)

# 检查点 + 早期停止
trainer = Trainer(
    callbacks=[
        ModelCheckpoint(monitor="val_loss", save_top_k=1),
        EarlyStopping(monitor="val_loss", patience=5),
    ],
)

# 日志
trainer = Trainer(
    logger=TensorBoardLogger("logs/", name="experiment"),
)

结语

PyTorch Lightning 是一个强大而灵活的深度学习框架，它通过消除样板代码、简化分布式训练、提供标准化接口，让研究者能够专注于模型创新，而将工程细节交给框架处理。

核心优势回顾：

✅ 代码简洁：减少 50%+ 的样板代码
✅ 易于扩展：一行代码从 CPU 扩展到多 GPU/多节点
✅ 内置最佳实践：混合精度、梯度累积、自动检查点等
✅ 生态系统完善：与 W&B、TensorBoard、MLflow 无缝集成
✅ 活跃的社区：持续更新和改进（2025 已发布 2.5+ 版本）

下一步行动：

动手实践：从 MNIST 示例开始，逐步尝试更复杂的项目
阅读源码：理解 Lightning 的内部实现原理
参与社区：在 GitHub 上提 Issue 或贡献代码
持续学习：关注官方博客和发布说明，了解最新特性

PyTorch Lightning 不仅仅是工具，更是一种思维方式——将科学研究从工程复杂性中解放出来，让创新更自由。无论你是研究者、工程师，还是学生，Lightning 都能帮助你更高效地实现深度学习项目。

最后提醒：深度学习框架不断演进，Lightning 也不例外。关注官方更新，定期查看新特性和最佳实践，保持学习和探索的态度。

祝你在 PyTorch Lightning 的旅程中收获满满！

参考资源

官方文档

教程与示例

社区

文章信息

作者: Sisyphus
发布日期: 2025-12-31
PyTorch Lightning 版本: 2.5+
适用人群: 深度学习研究者、工程师、学生

许可证: 本文内容遵循 CC BY-NC-SA 4.0 许可证。

Pytorch-lightning