非常游戏网
Python怎么保存和恢复PyTorch训练中断的断点_checkpoint打包模型与优化器状态

Python怎么保存和恢复PyTorch训练中断的断点_checkpoint打包模型与优化器状态

2026-05-28日常编程281838

真正可恢复的checkpoint至少需包含model.state_dict()、optimizer.state_dict()、当前epoch、最新loss,以及可选的lr_scheduler.state_dict()和随机数种子;漏掉optimizer.state_dict()会导致优化器状态丢失而收敛异常。

PyTorch保存checkpoint必须包含哪些关键字段

只保存模型参数(model.state_dict())不够,中断训练后恢复会报错或收敛异常。真正可恢复的checkpoint至少得打包:model.state_dict()optimizer.state_dict()、当前epoch、最新loss,以及可选的lr_scheduler.state_dict()和随机数种子。漏掉optimizer.state_dict()是新手最常踩的坑——模型能加载,但优化器内部动量、二阶矩等状态全丢,相当于从头开始训。

示例写法:

torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    'rng_state': torch.get_rng_state()
}, 'checkpoint.pth')

恢复checkpoint时model和optimizer要先初始化再load_state_dict

不能先load_state_dict()再定义模型或优化器——会报AttributeError: 'NoneType' object has no attribute 'load_state_dict'。必须确保modeloptimizer已实例化,且结构/参数名与保存时完全一致(比如用了nn.DataParallel但恢复时没加,或反之,会导致键不匹配)。

典型恢复流程:

Python 3.14.3

微软官方的 Python 扩展,是 VS Code 安装量最高的扩展(209M+)。集成 IntelliSense(通过 Pylance)、调试(通过 Python Debugger)、代码检查、格式化、重构和单元测试等功能。支持 Jupyter Notebook、虚拟环境管理和多 Python 版本切换。

  • 先构建好modeloptimizer(结构、参数名、device一致)
  • torch.load()读取checkpoint字典
  • 分别调用model.load_state_dict(checkpoint['model_state_dict'])optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
  • 手动恢复start_epoch = checkpoint['epoch'] + 1,别忘了+1,否则会重复训当前epoch

多GPU训练下state_dict键名不一致导致KeyError

nn.DataParallelDistributedDataParallel时,保存的model.state_dict()键名带module.前缀;单卡恢复时若模型没套DataParallel,就会报KeyError: 'module.conv1.weight'。解决方法不是硬改键名,而是统一加载逻辑:

  • 单卡恢复多卡保存的checkpoint:用{k.replace('module.', ''): v for k, v in state_dict.items()}预处理
  • 更稳妥做法:训练脚本开头就判断是否启用DP,并在保存/加载时保持一致封装
  • DDP场景下注意torch.load(..., map_location)指定device,否则默认加载到CPU,后续model.to(device)不自动迁移optimizer状态

保存路径、覆盖策略和文件损坏防护

别直接覆盖旧checkpoint——万一写入中途断电,新旧全丢。推荐用带epoch编号的命名(如checkpoint_epoch_42.pth),并配合shutil.copyfile()先写临时文件再原子重命名。另外,训练重启前务必检查checkpoint文件是否完整:用torch.load(..., weights_only=True)(PyTorch ≥2.3)可防恶意序列化攻击;小项目也可加简单校验,比如读取后立刻assert 'epoch' in checkpoint

容易被忽略的是:如果用了自定义学习率调度器(如继承torch.optim.lr_scheduler._LRScheduler),它的状态也得进checkpoint,否则step()调用会出错——这个点文档很少提,但实际debug花时间最多。