【Bug已解决】RuntimeError: Error(s) in loading state_dict for ResNet: 解决方案-创锋一号

【Bug已解决】RuntimeError: Error(s) in loading state_dict for ResNet: 解决方案

问题描述

在 PyTorch 中使用预训练模型或加载保存的模型权重时，RuntimeError: Error(s) in loading state_dict for ResNet是最常见的错误之一。这个错误发生在调用model.load_state_dict(torch.load('model.pth'))时，表示保存的权重文件（state_dict）与当前模型结构不匹配。

这个错误在以下场景中尤为常见：

使用预训练的 ResNet 模型并修改了网络结构（如修改最后一层以适应新的分类数）
在不同版本的 torchvision 之间迁移模型权重
从自定义训练的 ResNet 变体加载权重到标准 ResNet
使用 DataParallel 或 DistributedDataParallel 保存的模型加载到普通模型
模型定义代码与保存权重时的代码不一致

错误信息通常包含以下三种类型的不匹配：

Missing key(s)：state_dict 中缺少模型需要的某些层的权重
Unexpected key(s)：state_dict 中包含了模型不需要的额外层的权重
Size mismatch：某些层的权重形状不匹配

错误复现

场景一：修改了分类头导致 size mismatch

import torch import torchvision from torchvision.models import resnet50 # 加载预训练 ResNet50（1000 类 ImageNet） model_pretrained = resnet50(pretrained=True) torch.save(model_pretrained.state_dict(), 'resnet50_imagenet.pth') # 创建修改了最后一层的新模型（10 类） model = resnet50(num_classes=10) # 尝试加载预训练权重 model.load_state_dict(torch.load('resnet50_imagenet.pth'))

报错：

RuntimeError: Error(s) in loading state_dict for ResNet: size mismatch for fc.weight: copying a param with shape torch.Size([1000, 2048]) from checkpoint, the shape in current model is torch.Size([10, 2048]). size mismatch for fc.bias: copying a param with shape torch.Size([1000]) from checkpoint, the shape in current model is torch.Size([10]).

场景二：DataParallel 保存的模型加载到普通模型

# 保存时使用了 DataParallel model = resnet50() model = torch.nn.DataParallel(model) torch.save(model.state_dict(), 'resnet50_parallel.pth') # 加载到普通模型 model = resnet50() model.load_state_dict(torch.load('resnet50_parallel.pth'))

报错：

RuntimeError: Error(s) in loading state_dict for ResNet: Missing key(s) in state_dict: "conv1.weight", "bn1.weight", ... Unexpected key(s) in state_dict: "module.conv1.weight", "module.bn1.weight", ...

场景三：缺少某些层的权重

# 保存时只保存了部分权重 model = resnet50() partial_state = {k: v for k, v in model.state_dict().items() if 'fc' not in k} torch.save(partial_state, 'resnet50_partial.pth') # 加载完整模型 model = resnet50() model.load_state_dict(torch.load('resnet50_partial.pth'))

报错：

RuntimeError: Error(s) in loading state_dict for ResNet: Missing key(s) in state_dict: "fc.weight", "fc.bias"

根因分析

1. 分类头尺寸不匹配

ResNet 模型的最后一层（fc层）的输出维度取决于分类数。ImageNet 预训练模型有 1000 个类别，而自定义任务可能有不同数量的类别。直接加载会导致fc.weight和fc.bias的形状不匹配。

2. DataParallel/DistributedDataParallel 的`module.`前缀

当模型被DataParallel或DistributedDataParallel包装后，所有参数名会加上module.前缀。保存的 state_dict 中的键名如module.conv1.weight，而普通模型的键名是conv1.weight，导致键名不匹配。

3. 模型结构不一致

如果模型定义代码发生了变化（如添加/删除了层、修改了层的名称），保存的 state_dict 与当前模型结构不匹配。

4. torchvision 版本差异

不同版本的 torchvision 中，ResNet 的实现可能有细微差异（如层名、结构等），导致跨版本加载权重失败。

5. 自定义 ResNet 变体

如果使用了自定义的 ResNet 变体（如 ResNet + 注意力机制、修改的 Bottleneck 等），其 state_dict 的键名和形状与标准 ResNet 不同。

6. 严格模式（strict=True）

load_state_dict的默认参数strict=True要求 state_dict 的键与模型的键完全匹配。任何不匹配都会导致错误。

解决方案

方案一：使用`strict=False`加载部分权重

import torch import torchvision from torchvision.models import resnet50 # 创建修改了分类头的模型 model = resnet50(num_classes=10) # 加载预训练权重，使用 strict=False 忽略不匹配的层 state_dict = torch.load('resnet50_imagenet.pth') # 方法1：直接使用 strict=False（忽略 fc 层的不匹配） model.load_state_dict(state_dict, strict=False) # fc 层将使用随机初始化的权重 # 方法2：先删除不匹配的键，再加载 keys_to_remove = [k for k in state_dict if k.startswith('fc')] for k in keys_to_remove: del state_dict[k] model.load_state_dict(state_dict, strict=False)

方案二：手动替换分类头权重

import torch import torchvision from torchvision.models import resnet50 # 加载预训练 state_dict pretrained_dict = torch.load('resnet50_imagenet.pth') # 创建新模型 model = resnet50(num_classes=10) model_dict = model.state_dict() # 过滤掉 fc 层的权重 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape} # 更新模型 dict model_dict.update(pretrained_dict) model.load_state_dict(model_dict) print(f"Loaded {len(pretrained_dict)}/{len(model_dict)} layers")

方案三：处理 DataParallel 的`module.`前缀

import torch from torchvision.models import resnet50 # 加载 DataParallel 保存的权重 state_dict = torch.load('resnet50_parallel.pth') # 移除 module. 前缀 new_state_dict = {} for k, v in state_dict.items(): if k.startswith('module.'): new_state_dict[k[7:]] = v # 去掉 'module.'（7个字符） else: new_state_dict[k] = v # 加载到普通模型 model = resnet50() model.load_state_dict(new_state_dict)

方案四：使用 torchvision 的新 API 加载预训练权重

import torch import torchvision # 新版 API（torchvision >= 0.13） from torchvision.models import resnet50, ResNet50_Weights # 加载预训练权重 weights = ResNet50_Weights.DEFAULT model = resnet50(weights=weights) # 修改最后一层 num_features = model.fc.in_features model.fc = torch.nn.Linear(num_features, 10) # 10 类

完整修复代码

以下是一个完整的 ResNet 权重加载工具，处理各种不匹配情况：

""" ResNet state_dict 加载工具 处理各种常见的权重加载不匹配问题 """ import torch import torch.nn as nn from torchvision.models import resnet18, resnet34, resnet50, resnet101 from typing import Dict, Optional def load_resnet_with_transfer_learning( model: nn.Module, checkpoint_path: str, num_classes: int = None, strict: bool = False, verbose: bool = True ) -> nn.Module: """ 加载 ResNet 权重，处理分类头不匹配问题 Args: model: 目标模型 checkpoint_path: 权重文件路径 num_classes: 目标分类数（如果与 checkpoint 不同，fc 层将随机初始化） strict: 是否严格匹配 verbose: 是否打印详细信息 Returns: 加载了权重的模型 """ # 加载 checkpoint checkpoint = torch.load(checkpoint_path, map_location='cpu') # 处理 checkpoint 格式 if isinstance(checkpoint, dict): if 'state_dict' in checkpoint: state_dict = checkpoint['state_dict'] elif 'model_state_dict' in checkpoint: state_dict = checkpoint['model_state_dict'] else: state_dict = checkpoint else: state_dict = checkpoint # 清理键名 state_dict = clean_state_dict_keys(state_dict) # 获取模型当前的 state_dict model_dict = model.state_dict() # 分析不匹配 matched, mismatched_shapes, missing, unexpected = analyze_state_dict( state_dict, model_dict ) if verbose: print(f"权重加载分析:") print(f" 匹配的层: {len(matched)}") print(f" 形状不匹配: {len(mismatched_shapes)}") print(f" 缺失的层: {len(missing)}") print(f" 多余的层: {len(unexpected)}") if mismatched_shapes: print(f"\n 形状不匹配的层:") for k, (ckpt_shape, model_shape) in mismatched_shapes.items(): print(f" {k}: checkpoint={ckpt_shape}, model={model_shape}") if missing: print(f"\n 缺失的层（将使用随机初始化）:") for k in missing[:10]: print(f" {k}") if len(missing) > 10: print(f" ... 和其他 {len(missing) - 10} 个") # 过滤匹配的权重 filtered_dict = {} for k, v in state_dict.items(): if k in model_dict and v.shape == model_dict[k].shape: filtered_dict[k] = v # 加载 model_dict.update(filtered_dict) model.load_state_dict(model_dict, strict=strict) if verbose: loaded_count = len(filtered_dict) total_count = len(model_dict) print(f"\n成功加载 {loaded_count}/{total_count} 层权重") return model def clean_state_dict_keys(state_dict: Dict) -> Dict: """ 清理 state_dict 的键名 处理 DataParallel/DistributedDataParallel 的 module. 前缀 """ new_state_dict = {} for k, v in state_dict.items(): # 移除 module. 前缀 if k.startswith('module.'): k = k[7:] new_state_dict[k] = v return new_state_dict def analyze_state_dict( checkpoint_dict: Dict, model_dict: Dict ) -> tuple: """ 分析两个 state_dict 的差异 Returns: (matched, mismatched_shapes, missing, unexpected) """ matched = {} mismatched_shapes = {} missing = [] unexpected = [] for k, v in checkpoint_dict.items(): if k in model_dict: if v.shape == model_dict[k].shape: matched[k] = v else: mismatched_shapes[k] = (tuple(v.shape), tuple(model_dict[k].shape)) else: unexpected.append(k) for k in model_dict: if k not in checkpoint_dict: missing.append(k) return matched, mismatched_shapes, missing, unexpected def save_model_checkpoint( model: nn.Module, filepath: str, optimizer: Optional[torch.optim.Optimizer] = None, epoch: int = 0, **kwargs ): """ 保存模型 checkpoint（推荐格式） """ # 如果是 DataParallel，保存 module.xxx 而不是 module.module.xxx if isinstance(model, nn.DataParallel): state_dict = model.module.state_dict() elif isinstance(model, nn.parallel.DistributedDataParallel): state_dict = model.module.state_dict() else: state_dict = model.state_dict() checkpoint = { 'epoch': epoch, 'state_dict': state_dict, } if optimizer is not None: checkpoint['optimizer'] = optimizer.state_dict() checkpoint.update(kwargs) torch.save(checkpoint, filepath) print(f"Checkpoint saved to {filepath}") def load_model_checkpoint( model: nn.Module, filepath: str, optimizer: Optional[torch.optim.Optimizer] = None, device: str = 'cpu' ) -> dict: """ 加载模型 checkpoint """ checkpoint = torch.load(filepath, map_location=device) state_dict = checkpoint.get('state_dict', checkpoint.get('model_state_dict', checkpoint)) state_dict = clean_state_dict_keys(state_dict) model.load_state_dict(state_dict, strict=False) if optimizer is not None and 'optimizer' in checkpoint: optimizer.load_state_dict(checkpoint['optimizer']) return checkpoint # ============ 使用示例 ============ def demo_transfer_learning(): """演示迁移学习场景""" print("=" * 60) print("场景1: 迁移学习（修改分类头）") print("=" * 60) # 创建 10 分类的 ResNet50 model = resnet50(num_classes=10) # 模拟预训练权重（实际使用时替换为真实路径） pretrained_model = resnet50() # 默认 1000 类 torch.save(pretrained_model.state_dict(), 'fake_imagenet.pth') # 加载权重（自动处理 fc 层不匹配） model = load_resnet_with_transfer_learning( model, 'fake_imagenet.pth', num_classes=10, verbose=True ) # 验证 print(f"\n模型 fc 层: {model.fc}") print(f"fc.weight shape: {model.fc.weight.shape}") import os os.remove('fake_imagenet.pth') def demo_dataparallel(): """演示 DataParallel 权重加载""" print("\n" + "=" * 60) print("场景2: DataParallel 权重加载") print("=" * 60) # 模拟 DataParallel 保存的权重 model = resnet50() model = nn.DataParallel(model) torch.save(model.state_dict(), 'fake_parallel.pth') # 加载到普通模型 new_model = resnet50() new_model = load_resnet_with_transfer_learning( new_model, 'fake_parallel.pth', verbose=True ) import os os.remove('fake_parallel.pth') def demo_custom_resnet(): """演示自定义 ResNet 变体""" print("\n" + "=" * 60) print("场景3: 自定义 ResNet 变体") print("=" * 60) class CustomResNet(nn.Module): """自定义 ResNet，添加了额外的分类头""" def __init__(self, num_classes=10): super().__init__() self.backbone = resnet50() self.custom_head = nn.Linear(1000, num_classes) def forward(self, x): x = self.backbone(x) x = self.custom_head(x) return x # 创建自定义模型 custom_model = CustomResNet(num_classes=10) # 模拟预训练 ResNet50 权重 pretrained = resnet50() torch.save(pretrained.state_dict(), 'fake_pretrained.pth') # 加载（backbone 权重会匹配，custom_head 会随机初始化） checkpoint = torch.load('fake_pretrained.pth') checkpoint = clean_state_dict_keys(checkpoint) model_dict = custom_model.state_dict() matched = {k: v for k, v in checkpoint.items() if k in model_dict and v.shape == model_dict[k].shape} print(f"自定义模型总层数: {len(model_dict)}") print(f"匹配的层数: {len(matched)}") # 查看哪些层没有匹配 unmatched_model = [k for k in model_dict if k not in matched] print(f"未匹配的层: {unmatched_model[:5]}...") import os os.remove('fake_pretrained.pth') def demo_save_and_load(): """演示推荐的保存和加载方式""" print("\n" + "=" * 60) print("场景4: 推荐的保存/加载方式") print("=" * 60) model = resnet50(num_classes=10) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # 保存 save_model_checkpoint( model, 'demo_checkpoint.pth', optimizer=optimizer, epoch=10, best_acc=0.95 ) # 加载 new_model = resnet50(num_classes=10) new_optimizer = torch.optim.Adam(new_model.parameters(), lr=0.001) checkpoint = load_model_checkpoint( new_model, 'demo_checkpoint.pth', optimizer=new_optimizer ) print(f"恢复的 epoch: {checkpoint.get('epoch')}") print(f"恢复的 best_acc: {checkpoint.get('best_acc')}") import os os.remove('demo_checkpoint.pth') def main(): """主函数""" demo_transfer_learning() demo_dataparallel() demo_custom_resnet() demo_save_and_load() print("\n" + "=" * 60) print("所有演示完成！") print("=" * 60) if __name__ == '__main__': main()

常见陷阱与注意事项

1.`strict=False`的风险

使用strict=False会静默忽略不匹配的层。如果意外的不匹配（如模型结构错误）被忽略，可能导致模型性能下降且难以调试。建议在加载后检查哪些层被加载了，哪些被跳过了。

2. BN 层的`num_batches_tracked`

从 PyTorch 0.4 开始，BatchNorm 层增加了num_batches_tracked参数。在加载旧版本保存的权重时，可能缺少这个键。使用strict=False可以处理这种情况。

3. GPU 和 CPU 之间的加载

加载 GPU 上保存的权重到 CPU 模型时，需要指定map_location：

state_dict = torch.load('model.pth', map_location='cpu')

否则会报错：RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False.

4. 模型定义代码的版本控制

确保模型定义代码与保存权重时使用的是同一版本。如果模型定义发生了变化（如重命名了层），即使结构相同，键名不匹配也会导致加载失败。建议将模型定义代码与权重文件一起版本控制。

5. DataParallel vs DistributedDataParallel

DataParallel添加module.前缀，而DistributedDataParallel也添加module.前缀。但它们的处理方式略有不同。保存时建议总是保存model.module.state_dict()而不是model.state_dict()。

6. 检查权重的来源

在加载第三方提供的权重时，先检查其 state_dict 的键名：

state_dict = torch.load('unknown_model.pth') print(list(state_dict.keys())[:10]) # 查看前 10 个键名

这可以帮助你判断权重是来自标准模型还是 DataParallel 包装的模型。

7. 部分加载后需要初始化未加载的层

使用strict=False加载后，未匹配的层会保持随机初始化状态。确保这些层的初始化是合理的，或者在加载后手动初始化：

for m in model.modules(): if isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias)

8. torchvision 版本兼容性

不同版本的 torchvision 可能使用不同的预训练权重 URL 或权重格式。如果从旧版本迁移，建议使用新版本的 API 重新下载权重。

总结

RuntimeError: Error(s) in loading state_dict for ResNet错误虽然看起来复杂，但通常可以通过以下步骤解决：

分析不匹配原因：使用analyze_state_dict函数查看具体哪些层不匹配
处理module.前缀：使用clean_state_dict_keys移除 DataParallel 前缀
过滤不匹配的层：只加载形状匹配的权重
使用strict=False：允许部分加载
验证加载结果：检查哪些层成功加载，哪些被跳过

通过本文提供的完整工具代码，可以系统化地处理各种 ResNet 权重加载问题，使迁移学习和模型部署更加顺畅。记住，在修改模型结构后，始终验证权重加载的正确性，避免因静默跳过某些层而导致的性能问题。

企业官网建设流程全解析