📚 课程目录第2章 Checkpoint核心组件

2、Checkpoint核心组件：模型权重、优化器状态、学习率调度器、epoch与step信息、随机数生成器状态

很多人刚开始接触Checkpoint时，以为就是存个模型权重完事。说实话，我早年也这么干过。结果有一次训练跑了三天，机器挂了，恢复训练后发现loss曲线对不上，精度也差了一截。排查了半天，才发现是优化器状态丢了，学习率调度器也重置了。

那次之后，我彻底明白了——Checkpoint不是存个模型就完事了。它是一整套状态快照，缺一个组件，恢复训练就可能出问题。

核心观点：一个完整的Checkpoint，必须包含5个核心组件。少了任何一个，恢复训练时都会出现偏差。

2.1 模型权重——最基础，但也最容易出错

模型权重是Checkpoint的核心，这个大家都懂。但我想说的是，存什么、怎么存，其实有讲究。

我个人习惯用 model.state_dict() 来保存，而不是直接序列化整个模型对象。为什么？因为 state_dict 是一个有序字典，只包含可学习的参数。它更轻量，也更安全。你想想看，如果直接序列化整个模型，万一模型定义变了，加载时就可能报错。

# 推荐做法：只保存 state_dict
torch.save(model.state_dict(), 'checkpoint.pt')

# 恢复时
model = MyModel()
model.load_state_dict(torch.load('checkpoint.pt'))

小技巧：保存时加上 _use_new_zipfile_serialization=True，能显著减小文件体积，加载也更快。我在训练百亿参数模型时，这个参数帮我省了不少磁盘空间。

2.2 优化器状态——恢复训练的关键

优化器状态是什么？说白了，就是优化器内部维护的动量、方差等缓存数据。比如Adam优化器，它会记录每个参数的一阶动量和二阶动量。如果不保存这些，恢复训练时优化器就得从头开始积累，相当于白跑了好几个epoch。

我曾经在一个大模型训练任务中，因为没保存优化器状态，恢复后loss直接跳高了0.3，又花了整整两天才追回来。嗯，这个教训挺深刻的。

# 保存优化器状态
checkpoint = {
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(checkpoint, 'checkpoint.pt')

# 恢复时
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

注意：优化器状态和模型参数是绑定的。如果你改了模型结构（比如增减层数），旧的优化器状态就不能用了。加载时会报key不匹配的错误。

2.3 学习率调度器——容易被忽略的组件

学习率调度器保存的是当前的学习率、已调整次数、以及调度器内部的状态（比如CosineAnnealing的周期计数）。如果不保存，恢复训练后学习率会从初始值重新开始，这会导致训练震荡。

我见过一个案例：有人用ReduceLROnPlateau调度器，训练到第80个epoch时学习率降到了1e-5，结果机器挂了。恢复后没保存调度器状态，学习率又回到了1e-3，模型直接发散。

# 保存调度器状态
checkpoint['scheduler_state_dict'] = scheduler.state_dict()

# 恢复时
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

调度器类型	需要保存的状态	不保存的后果
StepLR	step_size, last_epoch	学习率调整时机错乱
CosineAnnealingLR	T_max, eta_min, last_epoch	余弦周期重置，学习率突变
ReduceLROnPlateau	cooldown_counter, best, num_bad_epochs	早停逻辑失效，学习率无法自适应

2.4 epoch与step信息——训练进度的坐标

这两个信息看似简单，但作用很大。epoch告诉你当前训练到第几轮，step告诉你当前batch的进度。恢复训练时，这两个值用于：

日志记录：保证训练曲线连续，不出现断点
数据加载：某些数据集需要按epoch切分，恢复时要从正确的epoch开始
验证逻辑：比如每5个epoch验证一次，恢复后要接着这个节奏

checkpoint['epoch'] = epoch
checkpoint['step'] = global_step

# 恢复时
start_epoch = checkpoint['epoch'] + 1
global_step = checkpoint['step']

注意：epoch和step的计数方式要统一。我习惯从0开始计数，这样恢复时直接 start_epoch = checkpoint['epoch'] + 1，逻辑清晰。

2.5 随机数生成器状态——最隐蔽的坑

这个组件，说实话，很多人压根没想过要保存。但它的影响非常大。

为什么会这样？因为训练过程中，很多操作依赖随机性：

数据加载时的shuffle
数据增强（随机裁剪、翻转、噪声）
Dropout的随机mask
初始化时的随机种子

如果不保存随机数生成器的状态，恢复训练后，数据增强的顺序会乱掉。比如原本第100个batch用了某种增强，恢复后可能变成另一种增强。虽然模型最终可能收敛，但训练过程不可复现，调试时非常头疼。

import torch
import random
import numpy as np

# 保存随机数生成器状态
checkpoint['torch_rng_state'] = torch.random.get_rng_state()
checkpoint['numpy_rng_state'] = np.random.get_state()
checkpoint['python_rng_state'] = random.getstate()

# 恢复时
torch.random.set_rng_state(checkpoint['torch_rng_state'])
np.random.set_state(checkpoint['numpy_rng_state'])
random.setstate(checkpoint['python_rng_state'])

避坑指南：如果你用了多个数据加载worker，每个worker有自己的随机种子。这种情况下，保存主进程的随机状态还不够，需要额外处理worker的随机状态。我曾经在这个问题上折腾了一整天，最后发现是DataLoader的worker_init_fn没处理好。

2.6 完整Checkpoint示例

好了，把上面5个组件整合起来，就是一个完整的Checkpoint。下面是我在实际项目中用的模板：

def save_checkpoint(model, optimizer, scheduler, epoch, step, 
                    torch_rng, numpy_rng, python_rng, filepath):
    checkpoint = {
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'epoch': epoch,
        'step': step,
        'torch_rng_state': torch_rng,
        'numpy_rng_state': numpy_rng,
        'python_rng_state': python_rng,
    }
    torch.save(checkpoint, filepath)
    print(f'Checkpoint saved at epoch {epoch}, step {step}')

def load_checkpoint(filepath, model, optimizer, scheduler):
    checkpoint = torch.load(filepath)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
    
    torch.random.set_rng_state(checkpoint['torch_rng_state'])
    np.random.set_state(checkpoint['numpy_rng_state'])
    random.setstate(checkpoint['python_rng_state'])
    
    return checkpoint['epoch'], checkpoint['step']

个人建议：除了这5个核心组件，我还会额外保存一些元信息，比如loss值、当前学习率、训练配置参数。这些信息在排查问题时很有用。你想想看，如果恢复后发现loss不对，至少能对比一下保存时的loss值，判断是训练问题还是加载问题。

总结一下这5个组件的关系：模型权重是骨架，优化器状态是肌肉记忆，学习率调度器是节奏控制器，epoch/step是坐标，随机数生成器是环境一致性保证。缺一个，恢复训练就不是无缝衔接，而是重新开始。