【Bug已解决】Problem with missing and unexpected keys while loading my model in Pytorch 解决方案问题描述在 PyTorch 中保存和加载模型时missing keys和unexpected keys是最常见的错误之一。当使用load_state_dict()加载模型权重时如果检查点中的键与当前模型的键不完全匹配PyTorch 会报告这些不匹配的键。典型的错误信息如下RuntimeError: Error(s) in loading state_dict for MyModel: Missing key(s) in state_dict: new_layer.weight, new_layer.bias Unexpected key(s) in state_dict: old_layer.weight, old_layer.bias或者在使用strictFalse时_IncompatibleKeys(missing_keys[new_layer.weight, new_layer.bias], unexpected_keys[old_layer.weight, old_layer.bias])这个问题在以下场景中频繁出现修改模型结构后加载旧检查点迁移学习中加载预训练权重不同版本模型之间的权重迁移使用 DataParallel/DistributedDataParallel 保存的模型错误复现错误一模型结构修改后加载旧检查点import torch import torch.nn as nn # 原始模型V1 class ModelV1(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 64, 3, padding1) self.conv2 nn.Conv2d(64, 128, 3, padding1) self.fc nn.Linear(128, 10) def forward(self, x): x torch.relu(self.conv1(x)) x torch.relu(self.conv2(x)) x x.view(x.size(0), -1) return self.fc(x[:, :128]) # 保存 V1 模型 model_v1 ModelV1() torch.save(model_v1.state_dict(), model_v1.pth) # 新模型V2- 添加了新层移除了旧层 class ModelV2(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 64, 3, padding1) self.conv2 nn.Conv2d(64, 128, 3, padding1) self.conv3 nn.Conv2d(128, 256, 3, padding1) # 新增层 self.fc nn.Linear(256, 10) # 修改了输入维度 # 移除了某些层 def forward(self, x): x torch.relu(self.conv1(x)) x torch.relu(self.conv2(x)) x torch.relu(self.conv3(x)) x x.view(x.size(0), -1) return self.fc(x[:, :256]) # 尝试加载 V1 的权重到 V2 model_v2 ModelV2() try: model_v2.load_state_dict(torch.load(model_v1.pth)) except RuntimeError as e: print(fError: {e}) # Missing key(s): conv3.weight, conv3.bias, fc.weight (shape mismatch) # Unexpected key(s): fc.weight, fc.bias (old shapes)错误二DataParallel 前缀问题import torch import torch.nn as nn # 使用 DataParallel 训练并保存 model nn.Sequential( nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 10) ) # 模拟 DataParallel if torch.cuda.device_count() 1: model nn.DataParallel(model) # 保存键名带有 module. 前缀 torch.save(model.state_dict(), dp_model.pth) # 加载到非 DataParallel 模型 new_model nn.Sequential( nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 10) ) checkpoint torch.load(dp_model.pth) print(Checkpoint keys:, list(checkpoint.keys())[:5]) # [module.0.weight, module.0.bias, module.2.weight, module.2.bias] try: new_model.load_state_dict(checkpoint) except RuntimeError as e: print(fError: {e}) # Missing keys: 0.weight, 0.bias, 2.weight, 2.bias # Unexpected keys: module.0.weight, module.0.bias, ...错误三strictFalse 仍然有问题# 使用 strictFalse result model_v2.load_state_dict( torch.load(model_v1.pth), strictFalse ) print(fMissing keys: {result.missing_keys}) print(fUnexpected keys: {result.unexpected_keys}) # 虽然不报错了但 missing keys 对应的层是随机初始化的 # 如果不处理这些层的权重是随机的模型性能会很差根因分析1. 模型结构变更当模型结构发生变化添加层、删除层、修改层名时state_dict 的键会相应变化。旧检查点的键无法完全匹配新模型。2. DataParallel/DistributedDataParallel 前缀使用nn.DataParallel或nn.DistributedDataParallel包装的模型其 state_dict 的所有键都会加上module.前缀。直接加载到未包装的模型会导致键名不匹配。3. 检查点格式不一致有些检查点只保存state_dict有些保存包含state_dict、optimizer、epoch等字段的字典。加载时需要正确提取state_dict。4. 自定义模块的命名在自定义模块中子模块的名称由属性名决定。如果重构代码时重命名了属性会导致键名变化。解决方案方案一键名映射和过滤import torch import torch.nn as nn from typing import Dict, List, Optional, Tuple def load_with_key_mapping( model: nn.Module, checkpoint: Dict[str, torch.Tensor], key_mapping: Optional[Dict[str, str]] None, skip_keys: Optional[List[str]] None, verbose: bool True ) - Tuple[List[str], List[str]]: 使用键名映射加载 state_dict Args: model: 目标模型 checkpoint: 检查点 state_dict key_mapping: 键名映射 {old_key: new_key} skip_keys: 要跳过的键 verbose: 是否打印详细信息 Returns: (missing_keys, unexpected_keys) if key_mapping is None: key_mapping {} if skip_keys is None: skip_keys [] model_sd model.state_dict() filtered {} for ckpt_key, ckpt_val in checkpoint.items(): # 跳过指定键 if ckpt_key in skip_keys: if verbose: print(f [SKIP] {ckpt_key}) continue # 应用键名映射 model_key key_mapping.get(ckpt_key, ckpt_key) # 去除 DataParallel 前缀 if model_key.startswith(module.): model_key model_key[7:] # 检查是否存在于模型中 if model_key not in model_sd: if verbose: print(f [UNEXPECTED] {ckpt_key} - {model_key} (not in model)) continue # 检查形状 if model_sd[model_key].shape ! ckpt_val.shape: if verbose: print(f [SHAPE MISMATCH] {model_key}: fmodel{model_sd[model_key].shape} vs ckpt{ckpt_val.shape}) continue filtered[model_key] ckpt_val if verbose: print(f [OK] {ckpt_key} - {model_key}) missing, unexpected model.load_state_dict(filtered, strictFalse) if verbose: if missing: print(f\n Missing keys: {missing}) if unexpected: print(f Unexpected keys: {unexpected}) print(f Loaded {len(filtered)} / {len(model_sd)} keys) return missing, unexpected方案二自动去除 DataParallel 前缀def strip_module_prefix(state_dict): 去除 DataParallel/DistributedDataParallel 的 module. 前缀 Args: state_dict: 原始 state_dict Returns: 清理后的 state_dict new_state_dict {} for key, value in state_dict.items(): new_key key if key.startswith(module.): new_key key[7:] # 去除 module. 前缀 new_state_dict[new_key] value return new_state_dict def load_model_safely( model: nn.Module, checkpoint_path: str, device: str cpu, strict: bool False, auto_strip_module: bool True, skip_mismatched: bool True ): 安全加载模型权重 Args: model: 目标模型 checkpoint_path: 检查点文件路径 device: 加载设备 strict: 是否严格匹配 auto_strip_module: 是否自动去除 module. 前缀 skip_mismatched: 是否跳过形状不匹配的键 Returns: (model, missing_keys, unexpected_keys) # 加载检查点 checkpoint torch.load(checkpoint_path, map_locationdevice, weights_onlyTrue) # 提取 state_dict if isinstance(checkpoint, dict) and state_dict in checkpoint: state_dict checkpoint[state_dict] elif isinstance(checkpoint, dict) and model_state_dict in checkpoint: state_dict checkpoint[model_state_dict] else: state_dict checkpoint # 去除 module. 前缀 if auto_strip_module: state_dict strip_module_prefix(state_dict) # 过滤形状不匹配的键 if skip_mismatched: model_sd model.state_dict() filtered {} for key, value in state_dict.items(): if key in model_sd: if model_sd[key].shape value.shape: filtered[key] value else: print(f [SKIP] Shape mismatch for {key}: fmodel{model_sd[key].shape} vs ckpt{value.shape}) else: print(f [SKIP] Key {key} not in model) state_dict filtered # 加载 missing, unexpected model.load_state_dict(state_dict, strictstrict) if missing: print(f\nMissing keys: {missing}) if unexpected: print(fUnexpected keys: {unexpected}) return model, missing, unexpected完整修复代码import torch import torch.nn as nn import torch.nn.functional as F from typing import Dict, List, Optional, Tuple, Any import copy import warnings class CheckpointLoader: 灵活的检查点加载器 功能 - 自动处理 DataParallel 前缀 - 支持键名映射 - 支持形状检查和跳过 - 支持部分加载 - 支持层名重命名 - 详细的加载报告 - 支持检查点格式检测 def __init__(self, model: nn.Module): self.model model self.model_state_dict model.state_dict() self.load_report { loaded: [], skipped: [], missing: [], unexpected: [], shape_mismatched: [], renamed: [], } def load( self, checkpoint: Any, strict: bool False, key_mapping: Optional[Dict[str, str]] None, skip_keys: Optional[List[str]] None, strip_module_prefix: bool True, shape_check: bool True, verbose: bool True, ) - Dict[str, List[str]]: 加载检查点 Args: checkpoint: 检查点state_dict 或包含 state_dict 的字典 strict: 是否严格匹配 key_mapping: 键名映射 {old_key: new_key} skip_keys: 要跳过的键名列表 strip_module_prefix: 是否去除 module. 前缀 shape_check: 是否检查形状 verbose: 是否打印详细信息 Returns: 加载报告字典 # 重置报告 self.load_report {k: [] for k in self.load_report} # 提取 state_dict state_dict self._extract_state_dict(checkpoint) # 去除 module. 前缀 if strip_module_prefix: state_dict self._strip_prefix(state_dict, module.) # 应用键名映射 if key_mapping: state_dict self._apply_key_mapping(state_dict, key_mapping) # 过滤 if skip_keys is None: skip_keys [] filtered_state_dict {} for key, value in state_dict.items(): if key in skip_keys: self.load_report[skipped].append(key) if verbose: print(f [SKIP] {key}) continue if key not in self.model_state_dict: self.load_report[unexpected].append(key) if verbose: print(f [UNEXPECTED] {key}) continue if shape_check: model_shape self.model_state_dict[key].shape ckpt_shape value.shape if model_shape ! ckpt_shape: self.load_report[shape_mismatched].append(key) if verbose: print(f [MISMATCH] {key}: fmodel{model_shape} vs ckpt{ckpt_shape}) continue filtered_state_dict[key] value self.load_report[loaded].append(key) # 加载 missing, unexpected self.model.load_state_dict( filtered_state_dict, strictFalse ) self.load_report[missing] list(missing) if strict: errors [] if self.load_report[missing]: errors.append(fMissing keys: {self.load_report[missing]}) if self.load_report[unexpected]: errors.append(fUnexpected keys: {self.load_report[unexpected]}) if self.load_report[shape_mismatched]: errors.append(fShape mismatched: {self.load_report[shape_mismatched]}) if errors: raise RuntimeError( Strict loading failed:\n \n.join(errors) ) if verbose: self._print_report() return self.load_report def load_from_file( self, path: str, device: str cpu, **kwargs ) - Dict[str, List[str]]: 从文件加载检查点 checkpoint torch.load(path, map_locationdevice, weights_onlyTrue) return self.load(checkpoint, **kwargs) def _extract_state_dict(self, checkpoint: Any) - Dict[str, torch.Tensor]: 从检查点中提取 state_dict if isinstance(checkpoint, dict): # 检查常见的键名 for key in [state_dict, model_state_dict, model, net]: if key in checkpoint: return checkpoint[key] # 如果没有找到假设整个字典就是 state_dict # 但先检查是否包含非张量值 if all(isinstance(v, torch.Tensor) for v in checkpoint.values()): return checkpoint else: warnings.warn( Checkpoint contains non-tensor values. Trying to use as state_dict directly. ) return {k: v for k, v in checkpoint.items() if isinstance(v, torch.Tensor)} return checkpoint def _strip_prefix(self, state_dict: Dict, prefix: str) - Dict: 去除键名前缀 new_sd {} for key, value in state_dict.items(): new_key key if key.startswith(prefix): new_key key[len(prefix):] self.load_report[renamed].append(f{key} - {new_key}) new_sd[new_key] value return new_sd def _apply_key_mapping(self, state_dict: Dict, mapping: Dict) - Dict: 应用键名映射 new_sd {} for key, value in state_dict.items(): new_key mapping.get(key, key) if new_key ! key: self.load_report[renamed].append(f{key} - {new_key}) new_sd[new_key] value return new_sd def _print_report(self): 打印加载报告 print(\n * 50) print(Checkpoint Loading Report) print( * 50) print(f Loaded: {len(self.load_report[loaded])} keys) print(f Missing: {len(self.load_report[missing])} keys) print(f Unexpected: {len(self.load_report[unexpected])} keys) print(f Shape mismatched:{len(self.load_report[shape_mismatched])} keys) print(f Skipped: {len(self.load_report[skipped])} keys) print(f Renamed: {len(self.load_report[renamed])} keys) if self.load_report[missing]: print(f\n Missing keys:) for k in self.load_report[missing]: print(f - {k}) if self.load_report[unexpected]: print(f\n Unexpected keys:) for k in self.load_report[unexpected]: print(f - {k}) if self.load_report[shape_mismatched]: print(f\n Shape mismatched keys:) for k in self.load_report[shape_mismatched]: print(f - {k}) if self.load_report[renamed]: print(f\n Renamed keys:) for k in self.load_report[renamed]: print(f - {k}) print( * 50) # 完整使用示例 class SimpleModel(nn.Module): def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d(1), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x class ModifiedModel(nn.Module): 修改后的模型添加了新层修改了分类器 def __init__(self, num_classes20): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.Conv2d(64, 128, 3, padding1), # 新增层 nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d(1), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128, 256), # 修改了维度 nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes), # 修改了类别数 ) def forward(self, x): x self.features(x) x self.classifier(x) return x def demo_missing_unexpected_keys(): 演示 missing/unexpected keys 问题及解决方案 print( * 60) print(Missing/Unexpected Keys 问题演示) print( * 60) # 创建并保存原始模型 print(\n[1] 创建并保存原始模型...) original_model SimpleModel(num_classes10) original_state_dict original_model.state_dict() print(f 原始模型键数量: {len(original_state_dict)}) # 模拟 DataParallel 前缀 dp_state_dict {} for key, value in original_state_dict.items(): dp_state_dict[fmodule.{key}] value print(f DataParallel 键数量: {len(dp_state_dict)}) # 创建修改后的模型 print(\n[2] 创建修改后的模型...) modified_model ModifiedModel(num_classes20) modified_state_dict modified_model.state_dict() print(f 修改后模型键数量: {len(modified_state_dict)}) # 场景1直接加载会失败 print(\n[3] 场景1直接加载strictTrue...) try: modified_model.load_state_dict(dp_state_dict, strictTrue) except RuntimeError as e: print(f Error: {str(e)[:200]}...) # 场景2使用 strictFalse print(\n[4] 场景2使用 strictFalse...) result modified_model.load_state_dict(dp_state_dict, strictFalse) print(f Missing keys: {len(result.missing_keys)}) print(f Unexpected keys: {len(result.unexpected_keys)}) # 场景3使用 CheckpointLoader print(\n[5] 场景3使用 CheckpointLoader...) modified_model2 ModifiedModel(num_classes20) loader CheckpointLoader(modified_model2) report loader.load(dp_state_dict, verboseTrue) # 场景4使用键名映射 print(\n[6] 场景4使用键名映射...) # 假设我们需要将 classifier.1 映射到 classifier.1不同维度 # 这里演示如何跳过不匹配的层 modified_model3 ModifiedModel(num_classes20) loader2 CheckpointLoader(modified_model3) report2 loader2.load( original_state_dict, skip_keys[ classifier.1.weight, classifier.1.bias, classifier.4.weight, classifier.4.bias, features.8.weight, features.8.bias, features.9.weight, features.9.bias, features.9.running_mean, features.9.running_var, features.9.num_batches_tracked, ], verboseTrue ) # 验证加载结果 print(\n[7] 验证加载结果...) # 检查共享层的权重是否正确加载 conv1_weight_original original_state_dict[features.0.weight] conv1_weight_loaded modified_model3.state_dict()[features.0.weight] is_loaded torch.equal(conv1_weight_original, conv1_weight_loaded) print(f features.0.weight 正确加载: {is_loaded}) # 检查新层是否为随机初始化 new_layer_weight modified_model3.state_dict()[features.8.weight] is_random not torch.equal(new_layer_weight, torch.zeros_like(new_layer_weight)) print(f features.8.weight 为随机初始化: {is_random}) print(\n * 60) print(演示完成!) print( * 60) def demo_save_load_best_practices(): 保存和加载的最佳实践 print(\n * 60) print(保存和加载的最佳实践) print( * 60) # 最佳实践保存完整的检查点 print(\n[1] 保存完整检查点...) model SimpleModel(num_classes10) optimizer torch.optim.Adam(model.parameters(), lr0.001) checkpoint { epoch: 10, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), loss: 0.05, model_info: { architecture: SimpleModel, num_classes: 10, pytorch_version: torch.__version__, }, } torch.save(checkpoint, best_checkpoint.pth) print( 保存了包含 epoch, model, optimizer, loss, model_info 的完整检查点) # 最佳实践加载完整检查点 print(\n[2] 加载完整检查点...) new_model SimpleModel(num_classes10) new_optimizer torch.optim.Adam(new_model.parameters(), lr0.001) loaded_ckpt torch.load(best_checkpoint.pth, weights_onlyTrue) loader CheckpointLoader(new_model) report loader.load( loaded_ckpt[model_state_dict], verboseFalse ) print(f 加载了 {len(report[loaded])} 层) # 恢复 optimizer 状态 new_optimizer.load_state_dict(loaded_ckpt[optimizer_state_dict]) print(f 恢复了 optimizer 状态) # 恢复 epoch start_epoch loaded_ckpt[epoch] 1 print(f 从 epoch {start_epoch} 继续训练) print(\n * 60) print(最佳实践演示完成!) print( * 60) if __name__ __main__: demo_missing_unexpected_keys() demo_save_load_best_practices()常见陷阱与注意事项1. strictFalse 不检查形状strictFalse只忽略键名不匹配不忽略形状不匹配。形状不匹配的键会导致错误必须手动过滤。2. module. 前缀的嵌套在嵌套使用 DataParallel 时可能出现module.module.双重前缀。需要循环去除。while key.startswith(module.): key key[7:]3. 检查点中的非张量数据有些检查点包含 Python 对象如学习率调度器状态使用weights_onlyTrue加载时会失败。需要根据检查点内容选择加载方式。4. 自定义层的 state_dict自定义层如果重写了state_dict()或load_state_dict()可能影响键名和加载行为。确保自定义实现与标准接口兼容。5. 模型.eval() 的必要性加载权重后如果模型包含 Dropout 或 BatchNorm必须调用model.eval()切换到推理模式。6. GPU 到 CPU 的设备迁移从 GPU 保存的检查点加载到 CPU 模型时使用map_locationcpu。7. 版本兼容性不同 PyTorch 版本的 state_dict 格式可能有细微差异。建议在检查点中记录 PyTorch 版本。总结missing keys和unexpected keys是 PyTorch 模型加载中的常见问题主要原因是模型结构变更、DataParallel 前缀和检查点格式差异。本文提供了完整的解决方案CheckpointLoader 工具类灵活处理各种加载场景支持键名映射、前缀去除、形状检查和跳过。自动前缀处理自动去除 DataParallel 的module.前缀。安全加载函数load_model_safely()提供一站式安全加载。最佳实践保存检查点时包含完整元信息便于后续恢复和调试。关键要点strictFalse只忽略键名不匹配不忽略形状不匹配DataParallel 的module.前缀是最常见的 unexpected keys 原因加载前检查检查点格式是否包含state_dict键保存检查点时附带模型架构和版本信息加载后调用model.eval()确保推理模式正确通过本文的工具和方法开发者可以轻松解决 PyTorch 中各种模型加载不匹配的问题实现灵活的迁移学习、模型恢复和版本迁移。