应保存 model.state_dict() 而非整个 model,因其仅含参数和缓冲区张量,轻量、可移植、加载快;完整保存需包含 epoch、optimizer 等元信息,并注意 eval() 模式、结构一致性和 device 映射。

为什么保存 model.state_dict() 而不是整个 model?
直接 torch.save(model, path) 会把模型类定义、构造函数参数、甚至 Python 对象引用(如自定义 forward 中的闭包、非 tensor 属性)一并序列化,体积大且跨环境易出错。而 model.state_dict() 只含 nn.Parameter 和已注册的 buffer(如 running_mean),纯张量字典,轻量、可移植、加载快。
保存时只存 state_dict 的标准写法
务必用 model.state_dict(),而非 model.module.state_dict()(除非用了 DataParallel 且当前是多卡训练后单卡加载);同时建议加 torch.save(..., _use_new_zipfile_serialization=True)(PyTorch ≥1.6 默认开启,但显式声明更稳妥):
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}, 'checkpoint.pth')
- 不推荐
torch.save(model.state_dict(), 'model.pth')—— 缺少元信息,后续无法知道训练轮次或优化器状态 - 若模型用了
torch.compile(),需先调用torch.compile(model, ...).uncompile()再取state_dict,否则可能报AttributeError: 'CompiledFunction' object has no attribute 'state_dict' - 保存前可调用
model.eval(),避免训练态 buffer(如 dropout mask)被意外存入
加载时必须重建模型结构再载入 state_dict
torch.load() 返回的是字典,不能直接当模型用。必须先实例化相同结构的模型对象,再用 load_state_dict():
model = MyModel() # 必须与保存时结构一致
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
- 若提示
Missing key(s) in state_dict:说明模型结构变了(比如新增了层),可用strict=False忽略缺失键,但要确认是否合理 - 若提示
Unexpected key(s) in state_dict:说明模型结构变精简了(如删了某层),同样可设strict=False,或检查是否误用了旧 checkpoint - 加载后别忘了
model.to(device),state_dict本身不含 device 信息
进一步压缩体积:用 torch.save(..., _pickle_protocol=4) 和半精度
默认 pickle 协议较旧,文件偏大;另外,若推理场景允许,可转为 float16 存储权重(注意 batch norm 的 running_var 需保持 float32 稳定性):
# 保存时转 half(仅适用于 inference)
state_dict = {k: v.half() if 'weight' in k or 'bias' in k else v
for k, v in model.state_dict().items()}
torch.save(state_dict, 'model_half.pth', _pickle_protocol=4)
-
_pickle_protocol=4在 Python 3.8+ 下生效,比默认的 protocol 2 小约 15–20% - 不要对
running_mean/running_var调用.half(),否则 BN 层在推理时数值不稳定 - 加载后需手动
model.half()并model.eval(),且确保所有输入也转为half
实际项目中,最常被忽略的是:保存时没确认模型是否处于 eval() 模式,导致 dropout 或 batch_norm 的临时状态混入;还有就是加载时直接 torch.load() 后当成模型调用,忘了重建结构——这两个点一旦出错,既不报错又结果异常,调试成本远高于体积本身。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











