1. PyTorch Lightning模型持久化入门指南

当你第一次听说PyTorch Lightning的模型持久化功能时，可能会觉得这不过是简单的保存和加载操作。但实际使用后你会发现，这里面藏着不少学问。我在去年做一个医疗影像分类项目时就深有体会——训练到第3天的模型因为服务器故障中断，幸好提前配置了ModelCheckpoint回调，才避免了从头开始的悲剧。

PyTorch Lightning的模型持久化主要解决三个核心问题：

• 训练中断恢复：自动保存最佳模型和最新模型，断电也不怕
• 实验管理：保存完整模型状态和超参数，方便横向对比不同实验
• 生产部署：将训练好的模型及其配置打包，便于后续推理使用

不同于原生PyTorch的torch.save()和torch.load()，PyTorch Lightning的持久化系统是围绕Trainer和LightningModule两个核心类构建的。这种设计最大的好处是，你不用再手动管理优化器状态、学习率调度器等训练细节，框架会自动帮你处理好这些"脏活"。

2. 自动保存：ModelCheckpoint回调实战

2.1 基础配置与监控策略

ModelCheckpoint是PyTorch Lightning中最强大的自动保存工具。先看一个我在实际项目中使用的配置示例：

from lightning.pytorch.callbacks import ModelCheckpoint

checkpoint_callback = ModelCheckpoint(
    dirpath="./checkpoints",
    filename="resnet-{epoch:02d}-{val_loss:.2f}",
    monitor="val_loss",
    mode="min",
    save_top_k=3,
    every_n_epochs=1,
    save_last=True
)

这个配置做了以下几件事：

1. 指定检查点保存在./checkpoints目录
2. 文件名包含epoch数和验证损失值
3. 监控val_loss指标并保留最小的3个检查点
4. 每1个epoch保存一次，同时始终保留最后一个epoch的模型

关键参数解析：

• monitor：要监控的指标名，通常用验证集的损失或准确率
• mode：取"min"（如监控loss）或"max"（如监控accuracy）
• save_top_k：保留表现最好的k个检查点，设为-1则保存所有
• save_last：是否额外保存最后一个epoch的模型（适合中断恢复）

2.2 高级保存策略与实战技巧

在实际项目中，你可能需要更复杂的保存策略。比如我在处理类别不平衡数据时，会同时监控多个指标：

class MultiMetricCheckpoint(ModelCheckpoint):
    def on_validation_end(self, trainer, pl_module):
        metrics = trainer.callback_metrics
        metrics["f1_score"] = 2*(metrics["precision"]*metrics["recall"])/(metrics["precision"]+metrics["recall"]+1e-8)
        super().on_validation_end(trainer, pl_module)

checkpoint_callback = MultiMetricCheckpoint(
    monitor="f1_score",
    mode="max",
    filename="best-f1-{epoch}-{f1_score:.3f}"
)

这个自定义回调在验证阶段计算F1分数，并以其作为保存依据。类似的，你还可以实现：

• 早停机制与检查点保存联动
• 根据学习率动态调整保存频率
• 在特定训练阶段（如loss陡降时）触发额外保存

常见踩坑点：

• 分布式训练时避免所有进程同时写文件（设置save_on_train_epoch_end=False）
• 监控指标必须在validation_step中通过self.log记录
• 大量检查点会占用磁盘空间，建议配合every_n_epochs使用

3. 手动保存的适用场景与陷阱

3.1 何时需要手动保存

虽然自动回调能满足大部分需求，但在以下场景手动保存更合适：

• 模型部署前：需要轻量化的推理专用模型
• 迁移学习：只保存骨干网络权重
• 临时快照：在特定实验步骤保存中间状态

PyTorch Lightning提供了两种手动保存方式：

# 方式1：保存完整训练状态（推荐）
trainer.save_checkpoint("full_checkpoint.ckpt")

# 方式2：仅保存模型权重（轻量但会丢失训练状态）
torch.save(model.state_dict(), "weights_only.pth")

3.2 手动保存的陷阱与解决方案

我在项目中最常遇到的坑是分布式训练时的死锁问题。当多个进程同时尝试保存模型时，可能会相互阻塞。解决方法很简单：

# 只在rank 0进程保存
if trainer.global_rank == 0:
    trainer.save_checkpoint("safe_save.ckpt")

另一个常见问题是手动保存的模型丢失超参数信息。这时可以用save_hyperparameters配合手动保存：

class MyModel(LightningModule):
    def __init__(self, lr=1e-3, hidden_dim=128):
        super().__init__()
        self.save_hyperparameters()
        # 模型定义...

# 保存时会自动包含超参数
trainer.save_checkpoint("with_hparams.ckpt")

4. 模型加载的完整工作流

4.1 标准加载方式

PyTorch Lightning推荐使用load_from_checkpoint加载模型：

# 加载完整模型（包含超参数）
model = MyLightningModule.load_from_checkpoint(
    checkpoint_path="path/to/checkpoint.ckpt"
)

# 加载后可以直接用于推理
outputs = model(torch.randn(1, 3, 224, 224))

关键点：

• 自动恢复模型结构和参数
• 保留原始训练时的超参数
• 支持GPU自动映射（与保存时设备无关）

4.2 高级加载技巧

在实际项目中，经常需要覆盖原始超参数。比如用预训练模型做迁移学习时：

model = MyLightningModule.load_from_checkpoint(
    checkpoint_path="pretrained.ckpt",
    lr=1e-4,  # 覆盖原始学习率
    hidden_dim=256  # 修改隐藏层维度
)

对于生产环境，你可能需要提取纯推理模型：

# 导出为TorchScript格式
script = model.to_torchscript()
torch.jit.save(script, "deployable_model.pt")

# 加载时不需要Lightning环境
deployed_model = torch.jit.load("deployable_model.pt")

5. 超参数管理的工程实践

5.1 save_hyperparameters深度解析

save_hyperparameters是PyTorch Lightning的超参数管理神器。它能自动捕获__init__参数：

class MyModel(LightningModule):
    def __init__(self, lr=1e-3, layers=4, dropout=0.1):
        super().__init__()
        self.save_hyperparameters()
        # 之后可以通过self.hparams访问参数
        self.layer1 = nn.Linear(10, self.hparams.layers)

高级用法：

• 手动指定要保存的参数：self.save_hyperparameters("lr", "dropout")
• 保存整个配置文件：self.save_hyperparameters(config_dict)
• 忽略特定参数：不将它们传入__init__

5.2 超参数版本控制

在大规模实验中，我推荐将超参数与模型检查点绑定：

checkpoint_callback = ModelCheckpoint(
    filename="{epoch}-{val_loss:.2f}-{hyperparams}",
    save_hyperparameters=True
)

这样每个检查点都会记录完整的超参数快照。你还可以用Lightning的TensorBoardLogger或MLFlowLogger实现更完善的实验追踪。

6. 生产环境下的最佳实践

经过多个项目的实战，我总结了以下经验：

1. 检查点命名规范：建议包含关键指标值（如val_acc=0.85）和实验标识
2. 定期清理策略：设置save_top_k配合磁盘监控脚本
3. 模型验证流程：加载后先用测试数据验证模型完整性
4. 元数据记录：在检查点中添加git commit hash等版本信息

一个典型的部署工作流如下：

# 训练阶段
trainer = Trainer(
    callbacks=[ModelCheckpoint(monitor="val_acc", mode="max")],
    logger=TensorBoardLogger("logs/")
)
trainer.fit(model)

# 部署阶段
best_model = MyModel.load_from_checkpoint(
    trainer.checkpoint_callback.best_model_path
)
torch.jit.save(best_model.to_torchscript(), "final_model.pt")

记住，好的持久化策略应该像时光机一样，能让你随时回到训练过程的任何关键节点。花时间配置好ModelCheckpoint回调，未来你会感谢现在的自己。