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

真正可恢复的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'。必须确保model和optimizer已实例化，且结构/参数名与保存时完全一致（比如用了nn.DataParallel但恢复时没加，或反之，会导致键不匹配）。

典型恢复流程：

Python 3.14.3

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

• 先构建好model和optimizer（结构、参数名、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.DataParallel或DistributedDataParallel时，保存的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花时间最多。