避坑指南：PyTorch模型保存时选torch.save还是state_dict？5个实际项目经验总结

PyTorch模型保存实战：从state_dict到完整模型的工程化选择

在深度学习项目部署和迭代过程中，模型保存就像程序员写代码不写注释——短期内看似省事，长期绝对是个灾难。作为PyTorch开发者，我们每天都在做选择题：该用torch.save(model)直接保存整个模型，还是老老实实用model.state_dict()只保存参数？这个看似简单的决策背后，藏着版本兼容性、团队协作效率、模型部署灵活性等一系列工程化考量。

1. 两种保存机制的本质差异

当我们把PyTorch模型保存到.pth文件时，实际上是在进行对象的序列化操作。但不同的保存方式会导致文件内容存在根本性差异：

# 完整模型保存示例 torch.save(model, 'full_model.pth') # 仅保存参数示例 torch.save(model.state_dict(), 'state_dict_only.pth')

完整模型保存会将以下内容打包进.pth文件：

• 模型类定义源代码的引用路径
• 所有可训练参数（权重和偏置）
• 模型结构定义（各层的连接方式）
• 前向传播方法的实现细节
• 自定义属性和辅助函数

而state_dict保存仅包含：

• 所有可训练参数的当前值
• 参数名称与张量的映射关系

关键提示：完整模型保存实际上会通过Python的pickle模块序列化整个模型对象，这可能导致在不同Python环境下出现兼容性问题。

下表对比了两种方式的核心特征：

2. 五种典型场景下的最佳实践

2.1 长期项目维护

在持续迭代的代码库中，state_dict是更可靠的选择。最近一个计算机视觉项目就踩了坑：团队用完整模型保存方式，半年后当需要调整模型结构时，发现原始类定义已被重构，导致历史模型全部无法加载。解决方法只能：

1. 找回旧版代码分支
2. 专门维护一个legacy.py存放废弃模型类
3. 额外编写转换脚本

# 糟糕的实践 - 强耦合于具体实现 class OldModel(nn.Module): ... # 好的实践 - 参数与结构解耦 new_model = NewModel() new_model.load_state_dict(torch.load('old_params.pth'))

2.2 跨框架部署

当需要将PyTorch模型部署到生产环境（如转换为ONNX格式）时，state_dict的优势更加明显。TensorRT等推理引擎通常需要：

1. 加载参数字典
2. 按需构建简化版推理模型
3. 进行格式转换

# ONNX转换示例 - 需要灵活控制模型结构 dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export(model, dummy_input, "model.onnx")

经验之谈：工业级部署往往需要去除训练专用的辅助层（如Dropout），此时state_dict方式可以自由重组模型结构。

2.3 学术研究共享

在论文复现或开源项目场景下，建议同时提供两种格式：

• 完整模型方便快速验证
• state_dict供高级用户灵活使用

例如HuggingFace模型库就采用这种双轨制：

model/ ├── pytorch_model.bin # state_dict └── config.json # 结构定义

2.4 迁移学习

进行模型微调时，不同保存策略会导致工作流差异：

完整模型流程：

1. 加载旧模型
2. 直接修改最后一层
3. 继续训练

state_dict流程：

1. 新建模型实例
2. 选择性加载参数
3. 冻结部分层
4. 修改输出层
5. 开始训练

# 迁移学习最佳实践 pretrained = torch.load('pretrained.pth') model = MyModel() model.load_state_dict(pretrained, strict=False) # 允许部分加载

2.5 多GPU训练部署

当使用DataParallel或DistributedDataParallel时，保存方式需要特别注意：

# 多GPU训练保存的正确姿势 model = nn.DataParallel(model) torch.save(model.module.state_dict(), 'multigpu.pth') # 注意.module

常见错误是直接保存包裹后的模型，会导致加载时出现意外的参数名前缀（如"module.conv1.weight"）。

3. 模型加载的七大陷阱与解决方案

3.1 版本不匹配报错

典型的错误信息：

AttributeError: Can't get attribute 'OldModel' on <module '__main__'>

解决方案：

1. 使用state_dict保存方式
2. 维护模型版本兼容层
3. 实现自定义加载逻辑：

def load_legacy_model(path): state_dict = torch.load(path, map_location='cpu') if 'state_dict' in state_dict: # 处理不同保存格式 state_dict = state_dict['state_dict'] # 处理参数名不匹配 new_state_dict = {} for k, v in state_dict.items(): name = k.replace('module.', '') # 去除多GPU前缀 new_state_dict[name] = v return new_state_dict

3.2 CUDA设备不匹配

当尝试将GPU保存的模型加载到CPU环境时：

RuntimeError: Attempting to deserialize object on CUDA device but torch.cuda.is_available() is False

正确处理方式：

# 指定加载设备 device = 'cuda' if torch.cuda.is_available() else 'cpu' state_dict = torch.load('model.pth', map_location=device)

3.3 参数形状不匹配

在修改模型结构后加载旧参数时常见：

RuntimeError: Error(s) in loading state_dict: size mismatch for fc.weight

调试checklist：

1. 打印新旧state_dict的键名对比
2. 使用strict=False参数部分加载
3. 手动过滤不匹配的参数

# 参数调试技巧 print("Current model keys:", model.state_dict().keys()) print("Loaded keys:", state_dict.keys()) # 选择性加载 model.load_state_dict(state_dict, strict=False)

4. 高级技巧：自定义保存策略

对于复杂项目，可以考虑混合保存策略：

# 自定义保存对象 checkpoint = { 'epoch': epoch, 'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict(), 'loss': loss, 'config': model_config # 保存必要的配置信息 } torch.save(checkpoint, 'checkpoint.pth')

恢复训练完整流程：

checkpoint = torch.load('checkpoint.pth') model = build_model(checkpoint['config']) # 根据配置重建 model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state'])

对于超大型模型，可以考虑分片保存：

# 参数分片保存 for name, param in model.named_parameters(): torch.save({name: param}, f'params/{name}.pt')

5. 性能优化与格式选择

.pth文件本质是Python的pickle格式，但还有其他选择：

二进制优化技巧：

# 使用最高效的pickle协议 torch.save(model, 'model.pth', pickle_protocol=5) # 启用压缩（Python 3.8+） torch.save(model, 'model.pth', pickle_protocol=5, _use_new_zipfile_serialization=True)

在部署到移动端时，可以考虑量化后再保存：

model = quantize_model(model) torch.save(model.state_dict(), 'quantized.pth')