只保存和加载 state_dict 是唯一能兼顾跨环境、可维护性与推理稳定性的做法,因其避免了 pickle 序列化导致的类定义、路径、版本等依赖问题,并需配合 map_location、手动构建模型、eval() 及明确命名规范。

只保存和加载 state_dict,不碰 torch.save(model, ...) —— 这是唯一能兼顾跨环境、可维护性与推理稳定性的做法。
为什么不能直接 torch.save(model, path)
这会把整个 Python 对象(含类定义、模块路径、临时属性)用 pickle 打包。一旦加载时缺少同名类、文件位置变了、PyTorch 版本不一致,就会报 AttributeError: Can't get attribute 'MyModel' on <module></module> 或更模糊的 Missing key 错误。这类问题往往只在部署或换机器时爆发,调试成本极高。
常见错误场景包括:
- 训练脚本里定义了
class Net(nn.Module): ...,但加载时没 import 它 - 模型用了
nn.DataParallel,保存时没取model.module.state_dict(),导致键名带module.前缀,单卡加载失败 - 保存时用了
torch.compile(model),但state_dict不包含编译图,加载后必须重新torch.compile()
torch.load() 必须配 map_location
GPU 上训练、CPU 上加载是最常见需求。不指定 map_location 会导致 RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False。
立即学习“Python免费学习笔记(深入)”;
正确写法是:
-
torch.load('model.pth', map_location='cpu')—— 强制加载到 CPU -
torch.load('model.pth', map_location=torch.device('cuda:0'))—— 指定 GPU 设备 -
torch.load('model.pth', map_location=lambda storage, loc: storage)—— 更底层的兼容写法,适用于旧版本或自定义设备映射
注意:map_location 必须在 torch.load() 里指定,而不是之后再调 model.to(device) —— 后者会先尝试把张量加载到原设备(比如 CUDA),再搬移,容易触发设备不可用错误。
加载前必须手动构造模型实例
state_dict 只是参数字典,它不包含任何网络结构信息。PyTorch 不提供“从权重反推模型”的能力。
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac,无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启:辅助功能 + 屏幕录制权限
所以加载流程只能是三步:
- 先写好模型类定义(哪怕只是复制粘贴一份)
- 用相同参数初始化模型:例如
model = MyModel(n_classes=10) - 再调
model.load_state_dict(torch.load(...))
漏掉第一步或参数不一致(比如 n_classes 改成 5),就会出现 size mismatch for fc.weight 错误。建议加载后加一句验证:
missing, unexpected = model.load_state_dict(state_dict, strict=False)
print(f"缺失参数: {missing}, 意外参数: {unexpected}")
如果 missing 非空,说明模型结构对不上;如果 unexpected 非空,可能是保存时混入了优化器状态,或模型用了 DataParallel 但没处理前缀。
推理模式与优化器状态要分开管理
用于推理时,加载完必须调 model.eval()。否则 BatchNorm 和 Dropout 仍按训练模式运行,输出不稳定甚至出错。
若需断点续训,优化器状态也得单独保存:
- 保存:
torch.save({'model': model.state_dict(), 'optimizer': optimizer.state_dict()}, 'checkpoint.pth') - 加载:
ckpt = torch.load('checkpoint.pth'); model.load_state_dict(ckpt['model']); optimizer.load_state_dict(ckpt['optimizer'])
切记:优化器的 state_dict 和模型的 state_dict 是两个独立字典,不能合并、不能互换、不能共用同一个文件名而不加区分。
最易被忽略的一点:.pth 或 .pt 文件本身不带元数据,你无法仅靠文件内容判断它存的是纯权重、完整模型还是 checkpoint —— 所以命名和文档必须明确,比如 model_weights.pth、checkpoint_epoch_10.pth,否则半年后你自己都可能搞混。

















