PyTorch-Lightning 中级教程：自定义模型检查点行为

廉咏燃

于 2025-06-01 09:01:49 发布

阅读量249

点赞数 3

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00676/article/details/148360222

PyTorch-Lightning 中级教程：自定义模型检查点行为

pytorch-lightning Lightning-AI/pytorch-lightning: PyTorch Lightning 是一个轻量级的高级接口，用于简化 PyTorch 中深度学习模型的训练流程。它抽象出了繁杂的工程细节，使研究者能够专注于模型本身的逻辑和实验设计，同时仍能充分利用PyTorch底层的灵活性。项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning

前言

在深度学习模型训练过程中，模型检查点（Checkpoint）是一个至关重要的功能。PyTorch-Lightning 提供了强大而灵活的检查点机制，可以帮助开发者高效地保存和恢复模型训练状态。本文将深入探讨如何自定义检查点行为，满足各种复杂场景的需求。

检查点基础概念

模型检查点通常包含以下内容：

模型权重参数
优化器状态
学习率调度器状态
训练进度信息（如当前epoch、step等）

PyTorch-Lightning 通过 ModelCheckpoint 回调提供了开箱即用的检查点功能，同时也支持高度自定义。

自定义检查点行为

1. 基本配置

最简单的检查点配置只需要几行代码：

from lightning.pytorch.callbacks import ModelCheckpoint

# 创建检查点回调
checkpoint_callback = ModelCheckpoint(
    dirpath="my/path/",  # 保存路径
    save_top_k=2,        # 保存最佳2个模型
    monitor="val_loss"   # 监控验证损失
)

# 将回调传递给Trainer
trainer = Trainer(callbacks=[checkpoint_callback])
trainer.fit(model)

# 获取最佳模型路径
best_model_path = checkpoint_callback.best_model_path

2. 监控自定义指标

你可以监控任何通过 self.log 记录的指标：

class LitModel(L.LightningModule):
    def training_step(self, batch, batch_idx):
        # 计算并记录自定义指标
        self.log("my_metric", x)

# 监控自定义指标
checkpoint_callback = ModelCheckpoint(monitor="my_metric")

高级检查点策略

1. 保存时机控制（When）

按训练步数保存：every_n_train_steps=N
按epoch间隔保存：every_n_epochs=N
按时间间隔保存：train_time_interval=timedelta(minutes=30)
训练epoch结束时保存：save_on_train_epoch_end=True

2. 保存内容选择（Which）

保存最新模型：save_last=True
保存最佳K个模型：save_top_k=10, monitor="val_loss", mode="min"
保存最近K个模型：save_top_k=10, monitor="global_step", mode="max"

# 保存最佳10个模型（基于val_loss）
checkpoint_callback = ModelCheckpoint(
    save_top_k=10,
    monitor="val_loss",
    mode="min",
    dirpath="my/path/",
    filename="model-{epoch:02d}-{val_loss:.2f}",
)

3. 保存内容精简（What）

默认情况下，检查点会保存完整训练状态。如果只需要模型权重：

checkpoint_callback = ModelCheckpoint(save_weights_only=True)

4. 保存路径定制（Where）

可以自定义文件名格式，包含监控指标：

checkpoint_callback = ModelCheckpoint(
    dirpath="my/path/",
    filename="model-{epoch:02d}-{val_loss:.2f}",
)

手动保存检查点

除了自动保存，你也可以手动保存和加载检查点：

# 训练后手动保存
trainer.fit(model)
trainer.save_checkpoint("manual_save.ckpt")

# 从检查点加载模型
new_model = MyLightningModule.load_from_checkpoint(
    checkpoint_path="manual_save.ckpt"
)

分布式训练注意事项

在分布式训练中，PyTorch-Lightning 会自动处理多进程保存问题：

trainer = Trainer(strategy="ddp")
model = MyLightningModule(hparams)
trainer.fit(model)

# 只在主进程保存，自动处理分布式策略
trainer.save_checkpoint("distributed.ckpt")

检查点模块化

检查点不仅可以保存模型状态，还可以保存数据模块和回调的状态：

# 保存数据模块状态
class MyDataModule(L.LightningDataModule):
    def on_save_checkpoint(self, checkpoint):
        checkpoint["data_state"] = self.data_state

# 保存回调状态
class MyCallback(L.Callback):
    def on_save_checkpoint(self, checkpoint):
        checkpoint["callback_state"] = self.callback_state

检查点修改钩子

可以在保存或加载检查点时修改其内容：

class LitModel(L.LightningModule):
    def on_save_checkpoint(self, checkpoint):
        # 添加自定义内容到检查点
        checkpoint["custom_data"] = self.custom_data

    def on_load_checkpoint(self, checkpoint):
        # 从检查点恢复自定义内容
        self.custom_data = checkpoint["custom_data"]

非严格加载检查点

默认情况下，检查点加载是严格的（参数名必须匹配）。可以禁用严格模式：

class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        # 禁用严格加载
        self.strict_loading = False
        
        # 预训练编码器（不更新）
        self.encoder = from_pretrained(...).requires_grad_(False)
        self.decoder = Decoder()

    def state_dict(self):
        # 只保存解码器状态
        return {k: v for k, v in super().state_dict().items() 
                if "encoder" not in k}