PyTorch 2.3 模型保存与加载：单卡/多卡场景下的4种模式与3个避坑点

PyTorch 2.3 模型保存与加载：单卡/多卡场景下的4种模式与3个避坑点

在深度学习项目开发中，模型保存与加载是连接训练与部署的关键环节。PyTorch作为当前最流行的深度学习框架之一，其模型序列化机制在实际工程应用中却存在不少"暗坑"，特别是在分布式训练与跨设备迁移场景下。本文将深入剖析PyTorch 2.3版本中模型保存与加载的四种核心场景，并揭示三个最常见的工程陷阱。

1. PyTorch模型保存基础：理解state_dict与序列化机制

PyTorch模型的本质是nn.Module类的实例，其核心参数通过state_dict()方法以有序字典的形式呈现。这个字典的键是各层的名称，值是对应的参数张量。理解这一点是掌握模型保存与加载的基础。

import torch from torchvision import models # 示例：查看ResNet模型的state_dict model = models.resnet18(pretrained=True) print(type(model.state_dict())) # <class 'collections.OrderedDict'> print(model.state_dict().keys()) # 输出所有层名称

PyTorch支持两种主要的保存方式：

1. 保存完整模型：包含模型结构和参数

torch.save(model, 'full_model.pt')

2. 仅保存参数（推荐方式）：

torch.save(model.state_dict(), 'model_weights.pt')

提示：生产环境中推荐只保存state_dict，这种方式更灵活且不受原始类定义变更的影响。

1.1 文件格式辨析：.pt、.pth与.pkl

实际上这三种格式在PyTorch中并无本质区别，选择.pt作为后缀是目前社区的最佳实践。

2. 单卡环境下的模型保存与加载

单GPU环境是最基础的场景，也是理解更复杂情况的基础。我们先看标准的保存与加载流程：

2.1 标准流程代码示例

# 保存模型 model = models.resnet50(pretrained=True) torch.save(model.state_dict(), 'resnet50_single_gpu.pt') # 加载模型 loaded_model = models.resnet50() # 必须保持相同模型结构 loaded_model.load_state_dict(torch.load('resnet50_single_gpu.pt')) loaded_model.eval() # 设置为评估模式

2.2 常见问题与验证方法

即使在这种简单场景下，也需要验证模型是否正确加载：

# 验证加载的模型与原模型参数一致 for (k1, v1), (k2, v2) in zip(model.state_dict().items(), loaded_model.state_dict().items()): assert k1 == k2 assert torch.allclose(v1, v2) print("验证通过！")

3. 多卡训练(DataParallel)下的模型处理

当使用nn.DataParallel进行多GPU训练时，模型会被自动包装，这带来了保存与加载的特殊性。

3.1 DataParallel的工作原理

model = nn.DataParallel(model) # 包装原始模型 print(next(model.parameters()).device) # 输出：cuda:0

DataParallel会在每个前向传播时将输入数据分割到各个GPU，最后汇总结果。这种设计导致模型参数名称前会自动添加"module."前缀。

3.2 多卡保存的三种处理策略

策略1：保存时去除DataParallel包装

torch.save(model.module.state_dict(), 'dp_model.pt') # 注意.module

策略2：加载时处理键名前缀

# 加载到单卡 state_dict = torch.load('dp_model.pt') state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} model.load_state_dict(state_dict)

策略3：加载到多卡环境

model = nn.DataParallel(model) # 重新包装 model.load_state_dict(torch.load('dp_model.pt')) # 保持原键名

注意：不同策略的选择取决于目标运行环境，策略1是最通用的做法。

4. 分布式训练(DistributedDataParallel)的特殊考量

nn.DistributedDataParallel（DDP）是PyTorch推荐的分布式训练方式，与DataParallel有显著区别。

4.1 DDP与DataParallel的关键差异

4.2 DDP模型的保存与加载

DDP模型也需要特殊处理，但与DataParallel类似：

# 保存时 if isinstance(model, nn.parallel.DistributedDataParallel): torch.save(model.module.state_dict(), 'ddp_model.pt') # 加载时 model = YourModelClass() model = nn.parallel.DistributedDataParallel(model) # 先包装 model.load_state_dict(torch.load('ddp_model.pt'))

5. 四种跨场景迁移模式详解

在实际工程中，我们经常需要在单卡与多卡环境之间迁移模型。以下是四种典型场景的处理方案。

5.1 单卡保存 → 单卡加载

这是最简单的情况，直接加载即可：

# 保存 torch.save(model.state_dict(), 'single_to_single.pt') # 加载 new_model = ModelClass() new_model.load_state_dict(torch.load('single_to_single.pt'))

5.2 单卡保存 → 多卡加载

需要先将模型加载到单卡，再包装为DataParallel：

# 加载到单卡 model = ModelClass() model.load_state_dict(torch.load('single_to_multi.pt')) # 包装为多卡模型 model = nn.DataParallel(model).cuda()

5.3 多卡保存 → 单卡加载

需要处理键名的"module."前缀：

state_dict = torch.load('multi_to_single.pt') state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} model.load_state_dict(state_dict)

5.4 多卡保存 → 多卡加载

保持键名一致即可：

model = nn.DataParallel(ModelClass()).cuda() model.load_state_dict(torch.load('multi_to_multi.pt'))

6. 三大工程陷阱与解决方案

在实际项目中，模型加载失败往往由以下几个常见问题导致。

6.1 设备不匹配问题

当尝试将GPU保存的模型加载到CPU环境时，会出现设备不匹配错误。

解决方案：使用map_location参数

# 强制加载到CPU torch.load('gpu_model.pt', map_location=torch.device('cpu')) # 自动选择可用设备 torch.load('model.pt', map_location=lambda storage, loc: storage)

6.2 键名不匹配问题

如前所述，DataParallel会引入"module."前缀，导致键名不一致。

诊断工具：

# 打印保存的和当前模型的键名差异 saved_keys = set(torch.load('model.pt').keys()) current_keys = set(model.state_dict().keys()) print("缺失的键:", saved_keys - current_keys) print("多余的键:", current_keys - saved_keys)

6.3 版本兼容性问题

不同PyTorch版本保存的模型可能存在兼容性问题。

最佳实践：

1. 在保存的模型文件名中包含PyTorch版本号
2. 使用相同的PyTorch版本进行训练和部署
3. 考虑导出为ONNX格式作为中间交换格式

7. 高级技巧与最佳实践

7.1 保存训练状态（checkpoint）

完整的训练状态保存应包括：

checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, 'lr_scheduler': scheduler.state_dict() } torch.save(checkpoint, 'checkpoint.pt')

7.2 模型瘦身技巧

对于部署场景，可以移除不必要的参数：

# 只保存可训练参数 pruned_dict = {k: v for k, v in model.state_dict().items() if v.requires_grad} torch.save(pruned_dict, 'pruned_model.pt')

7.3 安全加载外部模型

当加载不受信任的模型时，应采取安全措施：

# 使用pickle的安全加载方式 def safe_load(path): with open(path, 'rb') as f: return torch.load(f, map_location='cpu', pickle_module=pickle)

8. 性能优化与调试技巧

8.1 加速模型加载

对于大模型，可以启用共享内存：

torch.load('large_model.pt', map_location='cuda', mmap=True)

8.2 内存不足时的加载策略

# 分块加载大模型 def load_in_chunks(path, chunk_size=1024): state_dict = {} with open(path, 'rb') as f: for chunk in iter(lambda: f.read(chunk_size), b''): state_dict.update(torch.load(io.BytesIO(chunk))) return state_dict

8.3 模型加载性能对比

在实际项目中使用这些技术时，建议先在小规模数据上验证正确性，再应用到生产环境。PyTorch的模型序列化虽然简单易用，但细节决定成败，特别是在分布式和跨设备场景下。