PyTorch不推荐用pickle直接保存模型,因类定义和结构差异易致加载失败;应只保存state_dict,并同步保存optimizer、scheduler及epoch等状态以实现完整断点续训。

为什么直接保存 model 会报错或加载失败
PyTorch 不推荐用 pickle 直接序列化整个模型对象,因为模型类定义、模块结构、自定义方法等可能在不同环境不一致,导致 torch.load() 时反序列化失败。更关键的是,模型权重只是训练状态的一部分,optimizer、lr_scheduler、当前 epoch 和 loss 等也需要同步保存,否则“续训”只是重载权重,不是真正恢复训练上下文。
正确做法是只保存和加载 state_dict —— 它是 Python 字典,键为参数名(如 'conv1.weight'),值为 torch.Tensor,与模型结构解耦,安全、轻量、可移植。
保存完整训练状态:必须打包哪些 state_dict
断点续训不是只存模型权重,而是保存训练“快照”。典型需保存的有:
-
model.state_dict():模型各层参数 -
optimizer.state_dict():优化器内部状态(如 Adam 的一阶/二阶动量) -
lr_scheduler.state_dict()(如有):学习率调度器步数和当前 lr - 额外标量:
epoch、best_loss、global_step等,建议用普通字典字段存
示例保存代码:
立即学习“Python免费学习笔记(深入)”;
仅使用 Python 将 Claw Code 虚拟化框架与 OpenClaw 集成,提供 184 种工具和 200 多条命令,用于安全分析和代码质量...
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
'best_loss': best_loss,
}, 'checkpoint_epoch_{}.pt'.format(epoch))
加载时如何避免 Missing key 或 Unexpected key
这类错误本质是模型结构与保存时不一致:比如新增/删减了某层、改了层名、用了不同初始化方式。加载前必须确保模型实例已按原结构构建完成(包括所有 nn.Module 子模块),再调用 load_state_dict()。
常见应对策略:
- 加载时加
strict=False:跳过不匹配的 key(仅调试用,生产慎用) - 检查
model.state_dict().keys()和 checkpoint 中的 keys 是否一致 - 若模型含
nn.DataParallel,保存前用model.module.state_dict(),或加载时用torch.nn.DataParallel(model)包裹后再 load - 使用
torch.load(..., map_location=device)显式指定设备,避免 CPU/GPU 不匹配导致 tensor device error
加载后必须手动恢复训练上下文
调用 model.load_state_dict() 只恢复参数,不会自动设置 optimizer 的状态、学习率、当前 epoch。漏掉这些,等于从头开始训练,只是初始权重变了。
务必按顺序执行:
- 构建模型 →
model = MyModel() - 构建优化器 →
optimizer = torch.optim.Adam(model.parameters()) - 加载 checkpoint →
ckpt = torch.load('xxx.pt', map_location='cpu') - 分别加载 →
model.load_state_dict(ckpt['model_state_dict'])、optimizer.load_state_dict(ckpt['optimizer_state_dict']) - 恢复标量 →
start_epoch = ckpt['epoch'] + 1(注意 +1,因为刚加载完是上一轮结束状态) - 若用了
lr_scheduler,也需scheduler.load_state_dict(...)并注意其 step 逻辑是否依赖 epoch
最容易被忽略的是:model.train() 和 optimizer.zero_grad() 在续训循环开始前必须显式调用——加载不改变模型模式或梯度状态。

















