Checkpoint基础模型入门指南：从原理到实战部署

共计 2799 个字符，预计需要花费 7 分钟才能阅读完成。

为什么我们需要 checkpoint？

在训练深度模型时，经常会遇到各种意外情况导致训练中断：比如 GPU 显存不足、服务器宕机、或者人为误操作等。这时候如果没有保存中间状态，就只能从头开始训练，浪费了大量时间和计算资源。

• 典型场景：训练一个 ResNet 模型到第 20 个 epoch 时突然中断，所有进度丢失
• 直接损失：可能需要重新花费数小时甚至数天时间
• 隐性成本：难以复现中断前的模型状态，影响实验复现性

存储方案对比

PyTorch 提供了多种保存模型的方式，各有适用场景：

• torch.save()/torch.load()
• 优点：原生支持，能保存任意 PyTorch 对象（模型、张量等）

• 缺点：二进制格式，不易与其他框架交互

• HDF5 格式

• 优点：跨平台、可读性强，适合大模型参数存储

• 缺点：需要额外依赖库，不支持直接保存 PyTorch 对象

• ONNX 格式

• 优点：标准化模型交换格式
• 缺点：主要用于模型部署，不适合训练中间状态

对于训练过程中的 checkpoint，我们通常选择 torch.save()方案，因为它能完整保存训练上下文。

完整 checkpoint 实现

一个完整的 checkpoint 应该包含以下要素：

1. 模型当前参数（model.state_dict()）
2. 优化器状态（optimizer.state_dict()）
3. 当前 epoch 数
4. 其他训练状态（如学习率调度器）

# 保存 checkpoint 示例
checkpoint = {
    'epoch': current_epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'loss': current_loss,
}

torch.save(checkpoint, 'checkpoint_epoch_{}.pt'.format(current_epoch))

恢复训练时需要注意几个关键点：

1. 先重建模型和优化器的原始结构
2. 严格检查参数名称匹配（strict_load=True）
3. 恢复所有训练状态

# 加载 checkpoint 示例
checkpoint = torch.load('checkpoint_epoch_10.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1  # 从下一个 epoch 开始

分布式训练的特殊处理

在分布式数据并行 (DDP) 训练中，有几个易错点需要特别注意：

• 进程同步：只需在 rank= 0 的进程上保存 checkpoint

if torch.distributed.get_rank() == 0:
    torch.save(checkpoint, 'ddp_checkpoint.pt')

• 设备映射：加载时注意参数所在的设备

# 自动映射到当前设备
checkpoint = torch.load('ddp_checkpoint.pt', map_location='cuda:0')

• 屏障同步：保存前后需要同步所有进程

torch.distributed.barrier()

生产环境最佳实践

在实际项目中，我们还需要考虑以下工程化问题：

• 版本管理：建议采用包含 epoch 数和指标的命名方式

checkpoint_epoch=10_val_acc=0.85.pt

• 自动清理：只保留最近的 N 个 checkpoint

import glob
import os

# 保留最近 3 个 checkpoint
checkpoints = sorted(glob.glob('checkpoint_epoch_*.pt'))
for old_checkpoint in checkpoints[:-3]:
    os.remove(old_checkpoint)

• 校验机制：保存后立即验证文件完整性

def verify_checkpoint(path):
    try:
        checkpoint = torch.load(path)
        model.load_state_dict(checkpoint['model_state_dict'])
        return True
    except:
        return False

进阶思考

对于更复杂的生产系统，我们可以考虑：

1. 如何实现 checkpoint 的断点续传功能？
2. 在 Kubernetes 集群中如何设计自动恢复机制？
3. 如何优化超大模型的 checkpoint 存储效率？

动手实验

建议在 Colab 上尝试以下实验：

1. 用 MNIST 数据集训练一个简单 CNN
2. 在第 5 个 epoch 时手动中断训练
3. 从 checkpoint 恢复并继续训练
4. 比较恢复前后的模型表现

# 实验代码框架
import torch
import torchvision

model = CNN()
optimizer = torch.optim.Adam(model.parameters())
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5)

try:
    # 尝试加载 checkpoint
    checkpoint = torch.load('mnist_checkpoint.pt')
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch']
    print(f"从 epoch {start_epoch}恢复训练")
except:
    print("未找到 checkpoint，从头开始训练")
    start_epoch = 0

for epoch in range(start_epoch, 20):
    train_one_epoch(model, optimizer)

    # 保存 checkpoint
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),}, 'mnist_checkpoint.pt')

通过合理使用 checkpoint 机制，我们可以让模型训练过程更加健壮，有效避免意外中断带来的损失。建议在实践中形成规范的 checkpoint 管理流程，这对长期项目维护尤为重要。