1. 为什么只加载一部分比从头训练更常见刚接触深度学习那会儿我对预训练模型的理解特别朴素要么整个加载要么完全不加载。直到有次接了个医学影像分类的活儿拿ResNet50做骨干网络但官方权重是在ImageNet上训的最后那层全连接输出1000类我的任务只有3类。当时第一反应是把最后那层删掉重新建一个但中间那些层怎么办直接load_state_dict会报一堆key对不上的错整个模型加载直接失败。这个场景其实特别典型。实际项目里你拿到的预训练权重和你的模型结构几乎不可能完全一致原因无非这么几种分类数不同、骨干网络被裁剪过、加了自定义模块、或者只想复用特征提取部分而不要分类头。这时候就需要部分加载——只把能对上的参数搬过来对不上的随机初始化或者保留原样。很多人第一次遇到这个问题会去搜pytorch加载预训练模型报错然后看到一堆strictFalse的用法照抄上去发现不报错了但模型效果很差。问题出在哪strictFalse只是让加载不报错它不会告诉你哪些层被跳过了、哪些层被随机初始化了。如果你不小心把本该加载的层也跳过了模型等于半随机状态训练半天不收敛还找不到原因。所以这篇内容我想把部分加载这件事从头到尾讲透。不是只给你一行strictFalse就完事而是把背后的state_dict机制、key匹配逻辑、不同场景下的处理方案、以及我踩过的那些坑都摊开来说。适合已经能跑通完整训练流程、但在迁移学习时遇到结构不匹配问题的朋友。看完你应该能自己判断哪些层该加载、哪些该跳过、怎么验证加载结果是否符合预期。2. state_dict到底存了什么理解key匹配的底层逻辑2.1 state_dict不是模型本身是参数的字典很多人把state_dict和模型搞混。model.state_dict()返回的是一个Python字典严格说是OrderedDictkey是参数的名字value是对应的张量。它不包含模型结构信息只包含数值。这就解释了为什么加载时要求key能对上——因为它是按名字去匹配的。举个例子一个简单的两层网络import torch import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(10, 20) self.fc2 nn.Linear(20, 5) def forward(self, x): return self.fc2(self.fc1(x)) model SimpleNet() for k, v in model.state_dict().items(): print(k, v.shape)输出是fc1.weight torch.Size([20, 10]) fc1.bias torch.Size([20]) fc2.weight torch.Size([5, 20]) fc2.bias torch.Size([5])可以看到key就是模块名.参数名的格式。当你调用model.load_state_dict(pretrained_dict)时PyTorch会拿pretrained_dict里的每个key去model的state_dict里找同名key找到了就复制数值找不到就根据strict参数决定是报错还是跳过。2.2 strictTrue和strictFalse的真实区别默认strictTrue意思是两边key集合必须完全一致多一个少一个都报错。strictFalse则宽松很多只加载能对上的对不上的静默跳过。但这里有个关键细节很多人不知道——strictFalse时如果你模型里有某个层在预训练权重里不存在那个层会保持你初始化时的状态不会被重置。我见过有人以为strictFalse会把不匹配的层自动随机初始化其实不是。如果你的模型是先加载了别的权重再加载这个那些不匹配的层会保留上一次的数值。这个特性有时候是坑有时候是feature取决于你的流程设计。2.3 用missing_keys和unexpected_keys做诊断load_state_dict的返回值里有两个列表missing_keys和unexpected_keys。前者是你模型需要但预训练权重里没有的后者是预训练权重里有但你模型不需要的。这两个列表是排查问题的关键。pretrained_dict torch.load(pretrained.pth) model SimpleNet() result model.load_state_dict(pretrained_dict, strictFalse) print(缺失的key:, result.missing_keys) print(多余的key:, result.unexpected_keys)养成习惯每次部分加载后都打印这两个列表。如果missing_keys里出现了你本以为能加载的层说明key名字对不上可能是模块命名差异导致的。如果unexpected_keys一大堆说明预训练权重比你模型大很多可能你只用了骨干网络的一部分。提示missing_keys里的层如果是你故意要重新训练的比如分类头那没问题如果是你期望加载的层那就得检查命名了。3. 四种典型场景下的部分加载方案3.1 场景一分类数不同只换最后一层这是最常见的场景。预训练模型输出1000类你的任务输出10类。做法是先把模型结构改成你的分类数然后加载时过滤掉分类层的key。import torchvision.models as models import torch.nn as nn # 加载预训练模型 model models.resnet50(pretrainedTrue) # 替换最后的全连接层 num_features model.fc.in_features model.fc nn.Linear(num_features, 10) # 此时如果直接load_state_dict会报fc.weight和fc.bias的size不匹配 # 正确做法过滤掉fc层的key pretrained_dict models.resnet50(pretrainedTrue).state_dict() model_dict model.state_dict() # 方法一用字典推导式过滤 filtered_dict {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape model_dict[k].shape} model_dict.update(filtered_dict) model.load_state_dict(model_dict)这里我加了v.shape model_dict[k].shape这个条件因为光判断key存在还不够。如果key存在但形状不一样比如fc层直接update会报错。加上形状判断后fc层因为形状不匹配被自动过滤保持随机初始化状态。3.2 场景二只加载骨干网络自定义头全部重训有时候你不仅换了分类头还加了一些自定义模块比如注意力机制、多尺度融合层。这些层在预训练权重里根本不存在你只想加载骨干部分。class CustomModel(nn.Module): def __init__(self, backbone, num_classes): super().__init__() self.backbone backbone self.attention nn.Sequential( nn.Linear(2048, 512), nn.ReLU(), nn.Linear(512, 2048), nn.Sigmoid() ) self.classifier nn.Linear(2048, num_classes) def forward(self, x): features self.backbone(x) attn self.attention(features) return self.classifier(features * attn) backbone models.resnet50(pretrainedTrue) backbone.fc nn.Identity() # 去掉原分类头 model CustomModel(backbone, 10) # 加载时只取backbone相关的key pretrained_dict models.resnet50(pretrainedTrue).state_dict() model_dict model.state_dict() # 把预训练权重的key加上backbone.前缀来匹配 backbone_dict {fbackbone.{k}: v for k, v in pretrained_dict.items() if not k.startswith(fc.)} model_dict.update(backbone_dict) model.load_state_dict(model_dict)这个前缀映射的操作很关键。因为你的自定义模型把backbone包了一层state_dict的key会变成backbone.conv1.weight这种形式而原始预训练权重是conv1.weight。不做前缀处理一个key都对不上。3.3 场景三层数不一致只加载前N层有些时候你拿到的预训练模型比你需要的深比如用ResNet101的权重去初始化一个ResNet50。这时候只能加载前面能对上的层。pretrained_dict models.resnet101(pretrainedTrue).state_dict() model models.resnet50(pretrainedFalse) model_dict model.state_dict() # 逐层匹配只保留key和shape都一致的 matched_dict {} for k, v in pretrained_dict.items(): if k in model_dict and v.shape model_dict[k].shape: matched_dict[k] v print(f匹配成功: {k}) print(f总共匹配 {len(matched_dict)} 层模型共 {len(model_dict)} 层) model_dict.update(matched_dict) model.load_state_dict(model_dict)这种跨深度加载能匹配上的主要是前面的卷积层和BN层后面的layer3、layer4因为通道数和层数不同基本对不上。实际效果是浅层特征复用深层重新训练。对于数据量不大的任务这种策略往往比完全从头训练收敛快很多。3.4 场景四多GPU训练权重转单GPU用DataParallel或DDP训练保存的权重key前面会多一个module.前缀。加载到单GPU模型时全部对不上。# 多GPU保存的权重 pretrained_dict torch.load(multi_gpu_model.pth) # 去掉module.前缀 new_dict {} for k, v in pretrained_dict.items(): if k.startswith(module.): new_dict[k[7:]] v else: new_dict[k] v model.load_state_dict(new_dict)反过来单GPU权重加载到多GPU模型需要加前缀逻辑类似。这个坑几乎每个人都会踩一次记住module.这个前缀就行。4. 加载完了不等于对了验证与排查的完整链路4.1 加载后必须做的三项检查很多人加载完看到没报错就以为万事大吉结果训练时loss不降反升。我现在的习惯是加载后必做三个检查。第一打印missing_keys和unexpected_keys确认跳过的层符合预期。第二随机抽几个关键层的参数对比加载前后的数值是否真的变了。第三用同一批数据跑一次前向传播看输出是否正常不是NaN量级合理。# 检查一key诊断 result model.load_state_dict(filtered_dict, strictFalse) print(缺失:, result.missing_keys) print(多余:, result.unexpected_keys) # 检查二参数数值对比 for name, param in model.named_parameters(): if name in filtered_dict: loaded filtered_dict[name] if torch.equal(param.data, loaded): print(f{name} 加载成功) else: print(f{name} 数值不一致) break # 检查三前向传播 model.eval() with torch.no_grad(): dummy torch.randn(2, 3, 224, 224) out model(dummy) print(输出形状:, out.shape) print(输出范围:, out.min().item(), out.max().item()) print(是否有NaN:, torch.isnan(out).any().item())4.2 一个真实的排查案例之前有个项目我用EfficientNet-B0做骨干加载官方权重后训练前几个epoch loss正常下降到第5个epoch突然爆炸。排查了半天最后发现问题出在BN层的num_batches_tracked这个buffer上。EfficientNet的预训练权重里BN层的num_batches_tracked是一个很大的数因为在大数据集上训了很久而我的模型初始化时它是0。加载时这个buffer也被加载了导致BN层在计算时使用了预训练阶段的统计量和我的数据分布不匹配。前几个epoch因为momentum更新还没覆盖掉到后面就发散了。解决方案是在加载时过滤掉所有num_batches_trackedfiltered_dict {k: v for k, v in pretrained_dict.items() if num_batches_tracked not in k}这个坑很隐蔽因为num_batches_tracked不是可学习参数打印模型参数时看不到它但它在state_dict里确实存在。如果你用strictFalse它会被静默加载然后埋下隐患。4.3 学习率策略要跟着调整部分加载后不同层的新鲜程度是不一样的。加载了预训练权重的层已经学到了通用特征而随机初始化的层还是白纸。如果对所有层用同一个学习率要么预训练层被破坏要么新层学得太慢。我的做法是分组设置学习率预训练部分用小学习率比如1e-4新初始化的部分用大学习率比如1e-3。在PyTorch里通过参数组实现pretrained_params [] new_params [] for name, param in model.named_parameters(): if name in filtered_dict: pretrained_params.append(param) else: new_params.append(param) optimizer torch.optim.Adam([ {params: pretrained_params, lr: 1e-4}, {params: new_params, lr: 1e-3} ])如果数据量特别小还可以考虑先冻结预训练层只训新层几个epoch再解冻一起微调。这个策略在少样本场景下效果很明显。5. 那些年我踩过的加载坑与应对技巧5.1 坑一key名字差一个字母全部对不上有次用某个开源实现的骨干网络它的模块命名是features.0.weight而我用的官方实现是conv1.weight。结构完全一样就是命名不同导致一个key都加载不上。当时没打印missing_keys训练了半天发现效果和随机初始化一样才回头排查。应对方法加载前先打印两边state_dict的前几个key对比一下。如果发现命名体系不同写一个映射字典手动转换。# 查看两边的key print(预训练权重keys:, list(pretrained_dict.keys())[:5]) print(模型keys:, list(model.state_dict().keys())[:5]) # 如果命名不同建立映射 name_mapping { features.0.weight: conv1.weight, features.0.bias: conv1.bias, # ... } remapped_dict {name_mapping.get(k, k): v for k, v in pretrained_dict.items()}5.2 坑二DataParallel保存的权重加载到CPU模型在GPU上用DataParallel训练保存的权重加载到CPU上推理时除了module.前缀问题还会遇到设备不匹配。torch.load默认会把权重加载到保存时的设备上如果保存时在GPU而你现在只有CPU会报错。# 加载时指定map_location pretrained_dict torch.load(model.pth, map_locationcpu) # 或者加载后统一处理 pretrained_dict torch.load(model.pth, map_locationlambda storage, loc: storage)map_location这个参数建议养成习惯每次都写不管当前有没有GPU。这样代码在不同机器上迁移时不会出问题。5.3 坑三加载了权重但忘了切换train/eval模式这个坑和加载本身无关但经常一起出现。加载预训练权重后如果模型里有Dropout和BatchNorm忘记调用model.train()或model.eval()会导致行为异常。特别是做特征提取时必须model.eval()并且torch.no_grad()否则BN层会更新统计量Dropout会随机丢弃提取的特征不稳定。# 特征提取模式 model.eval() with torch.no_grad(): features model(input_tensor) # 微调模式 model.train() output model(input_tensor)5.4 坑四部分加载后直接保存下次加载又出问题有时候你部分加载后想保存这个半初始化的模型方便下次直接用。但如果保存的是整个模型对象torch.save(model)下次加载时如果代码里的类定义变了会反序列化失败。正确做法是只保存state_dict。# 推荐只保存state_dict torch.save(model.state_dict(), partially_loaded.pth) # 加载时先建模型再加载 model CustomModel(...) model.load_state_dict(torch.load(partially_loaded.pth))这样模型结构和权重分离代码重构时不会因为类定义变化导致加载失败。5.5 一个实用的加载工具函数踩了这么多坑之后我封装了一个通用的部分加载函数每次直接用省得重复写过滤逻辑def load_pretrained_partial(model, pretrained_path, skip_keysNone, verboseTrue, map_locationcpu): 部分加载预训练权重 skip_keys: 需要跳过的key前缀列表如[fc, classifier] pretrained_dict torch.load(pretrained_path, map_locationmap_location) model_dict model.state_dict() # 去掉module.前缀 new_dict {} for k, v in pretrained_dict.items(): if k.startswith(module.): k k[7:] new_dict[k] v pretrained_dict new_dict # 过滤skip_keys if skip_keys: pretrained_dict {k: v for k, v in pretrained_dict.items() if not any(k.startswith(sk) for sk in skip_keys)} # 过滤num_batches_tracked pretrained_dict {k: v for k, v in pretrained_dict.items() if num_batches_tracked not in k} # 只保留key和shape都匹配的 matched_dict {} for k, v in pretrained_dict.items(): if k in model_dict and v.shape model_dict[k].shape: matched_dict[k] v model_dict.update(matched_dict) result model.load_state_dict(model_dict, strictFalse) if verbose: print(f成功加载 {len(matched_dict)}/{len(model_dict)} 层) if result.missing_keys: print(f缺失的key: {result.missing_keys}) return model, result这个函数覆盖了前缀处理、skip过滤、形状匹配、buffer过滤几个常见需求。实际用的时候根据情况传参就行。6. 不同框架版本下的行为差异PyTorch不同版本在load_state_dict的行为上有一些细微差别不注意的话升级版本后可能出问题。在1.0到1.4版本之间strictFalse时如果key存在但shape不匹配会直接报错而不是跳过。从1.5开始shape不匹配的key会被归入unexpected_keys并跳过。所以如果你在旧版本上写的代码依赖了shape不匹配会报错这个行为来做判断升级后逻辑就失效了。另外assign参数是1.13才加入的。默认assignFalse时加载是原地复制数值assignTrue时会用预训练权重的张量替换模型原有的张量。区别在于assignTrue会保留预训练权重的设备、dtype等属性。如果你从GPU权重加载到CPU模型用assignTrue可能会把模型参数变成CUDA张量导致后续CPU计算报错。# 安全做法明确指定map_location不用assign model.load_state_dict(torch.load(path, map_locationcpu), strictFalse)还有一个容易忽略的点torch.load在2.6版本之后默认weights_onlyTrue这意味着它只能加载纯张量字典不能加载包含自定义类的完整模型。如果你保存的是整个模型对象加载时会报错。解决方案是保存时就用state_dict或者加载时显式传weights_onlyFalse但这样有安全风险只加载可信来源的文件。7. 从部分加载延伸到迁移学习的策略选择部分加载只是手段最终目的是让模型在新任务上表现好。加载完之后怎么训其实比怎么加载更影响结果。如果你的新任务数据和预训练数据分布接近比如都是自然图像那加载大部分层、只换分类头、用小学习率微调就够了。如果分布差异大比如自然图像到医学影像浅层特征可能也不适用这时候可以考虑只加载前几层后面全部重训或者干脆从头训练但用预训练权重做初始化。数据量也是个关键因素。数据少的时候冻结预训练层只训新层能防止过拟合数据多的时候全部解冻一起微调效果更好。我一般的做法是先用小学习率全部微调跑几个epoch看验证集表现如果过拟合就冻结浅层如果欠拟合就加大学习率或解冻更多层。还有一个技巧是分层设置学习率衰减浅层学习率最小深层逐渐增大。因为浅层学到的是通用特征边缘、纹理不需要大改深层学到的是任务相关特征需要更多调整。# 分层学习率示例 params [ {params: model.conv1.parameters(), lr: 1e-5}, {params: model.layer1.parameters(), lr: 1e-5}, {params: model.layer2.parameters(), lr: 5e-5}, {params: model.layer3.parameters(), lr: 1e-4}, {params: model.layer4.parameters(), lr: 5e-4}, {params: model.fc.parameters(), lr: 1e-3}, ] optimizer torch.optim.Adam(params)这套组合拳打下来大部分迁移学习场景都能覆盖。核心思路就是加载时精细控制哪些层复用训练时区别对待不同层的更新幅度。我在实际项目里最深的体会是部分加载这件事写代码只占两成时间剩下八成都在验证和调试。每次加载完多花五分钟做检查能省下后面几小时的排查。尤其是missing_keys和unexpected_keys这两个列表一定要看不要嫌麻烦。 
长河君 · 2026/10/14 0:20:07 · 阅读 1.2k PyTorch预训练模型部分加载：state_dict机制与迁移学习实战
1. 为什么只加载一部分比从头训练更常见刚接触深度学习那会儿我对预训练模型的理解特别朴素要么整个加载要么完全不加载。直到有次接了个医学影像分类的活儿拿ResNet50做骨干网络但官方权重是在ImageNet上训的最后那层全连接输出1000类我的任务只有3类。当时第一反应是把最后那层删掉重新建一个但中间那些层怎么办直接load_state_dict会报一堆key对不上的错整个模型加载直接失败。这个场景其实特别典型。实际项目里你拿到的预训练权重和你的模型结构几乎不可能完全一致原因无非这么几种分类数不同、骨干网络被裁剪过、加了自定义模块、或者只想复用特征提取部分而不要分类头。这时候就需要部分加载——只把能对上的参数搬过来对不上的随机初始化或者保留原样。很多人第一次遇到这个问题会去搜pytorch加载预训练模型报错然后看到一堆strictFalse的用法照抄上去发现不报错了但模型效果很差。问题出在哪strictFalse只是让加载不报错它不会告诉你哪些层被跳过了、哪些层被随机初始化了。如果你不小心把本该加载的层也跳过了模型等于半随机状态训练半天不收敛还找不到原因。所以这篇内容我想把部分加载这件事从头到尾讲透。不是只给你一行strictFalse就完事而是把背后的state_dict机制、key匹配逻辑、不同场景下的处理方案、以及我踩过的那些坑都摊开来说。适合已经能跑通完整训练流程、但在迁移学习时遇到结构不匹配问题的朋友。看完你应该能自己判断哪些层该加载、哪些该跳过、怎么验证加载结果是否符合预期。2. state_dict到底存了什么理解key匹配的底层逻辑2.1 state_dict不是模型本身是参数的字典很多人把state_dict和模型搞混。model.state_dict()返回的是一个Python字典严格说是OrderedDictkey是参数的名字value是对应的张量。它不包含模型结构信息只包含数值。这就解释了为什么加载时要求key能对上——因为它是按名字去匹配的。举个例子一个简单的两层网络import torch import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(10, 20) self.fc2 nn.Linear(20, 5) def forward(self, x): return self.fc2(self.fc1(x)) model SimpleNet() for k, v in model.state_dict().items(): print(k, v.shape)输出是fc1.weight torch.Size([20, 10]) fc1.bias torch.Size([20]) fc2.weight torch.Size([5, 20]) fc2.bias torch.Size([5])可以看到key就是模块名.参数名的格式。当你调用model.load_state_dict(pretrained_dict)时PyTorch会拿pretrained_dict里的每个key去model的state_dict里找同名key找到了就复制数值找不到就根据strict参数决定是报错还是跳过。2.2 strictTrue和strictFalse的真实区别默认strictTrue意思是两边key集合必须完全一致多一个少一个都报错。strictFalse则宽松很多只加载能对上的对不上的静默跳过。但这里有个关键细节很多人不知道——strictFalse时如果你模型里有某个层在预训练权重里不存在那个层会保持你初始化时的状态不会被重置。我见过有人以为strictFalse会把不匹配的层自动随机初始化其实不是。如果你的模型是先加载了别的权重再加载这个那些不匹配的层会保留上一次的数值。这个特性有时候是坑有时候是feature取决于你的流程设计。2.3 用missing_keys和unexpected_keys做诊断load_state_dict的返回值里有两个列表missing_keys和unexpected_keys。前者是你模型需要但预训练权重里没有的后者是预训练权重里有但你模型不需要的。这两个列表是排查问题的关键。pretrained_dict torch.load(pretrained.pth) model SimpleNet() result model.load_state_dict(pretrained_dict, strictFalse) print(缺失的key:, result.missing_keys) print(多余的key:, result.unexpected_keys)养成习惯每次部分加载后都打印这两个列表。如果missing_keys里出现了你本以为能加载的层说明key名字对不上可能是模块命名差异导致的。如果unexpected_keys一大堆说明预训练权重比你模型大很多可能你只用了骨干网络的一部分。提示missing_keys里的层如果是你故意要重新训练的比如分类头那没问题如果是你期望加载的层那就得检查命名了。3. 四种典型场景下的部分加载方案3.1 场景一分类数不同只换最后一层这是最常见的场景。预训练模型输出1000类你的任务输出10类。做法是先把模型结构改成你的分类数然后加载时过滤掉分类层的key。import torchvision.models as models import torch.nn as nn # 加载预训练模型 model models.resnet50(pretrainedTrue) # 替换最后的全连接层 num_features model.fc.in_features model.fc nn.Linear(num_features, 10) # 此时如果直接load_state_dict会报fc.weight和fc.bias的size不匹配 # 正确做法过滤掉fc层的key pretrained_dict models.resnet50(pretrainedTrue).state_dict() model_dict model.state_dict() # 方法一用字典推导式过滤 filtered_dict {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape model_dict[k].shape} model_dict.update(filtered_dict) model.load_state_dict(model_dict)这里我加了v.shape model_dict[k].shape这个条件因为光判断key存在还不够。如果key存在但形状不一样比如fc层直接update会报错。加上形状判断后fc层因为形状不匹配被自动过滤保持随机初始化状态。3.2 场景二只加载骨干网络自定义头全部重训有时候你不仅换了分类头还加了一些自定义模块比如注意力机制、多尺度融合层。这些层在预训练权重里根本不存在你只想加载骨干部分。class CustomModel(nn.Module): def __init__(self, backbone, num_classes): super().__init__() self.backbone backbone self.attention nn.Sequential( nn.Linear(2048, 512), nn.ReLU(), nn.Linear(512, 2048), nn.Sigmoid() ) self.classifier nn.Linear(2048, num_classes) def forward(self, x): features self.backbone(x) attn self.attention(features) return self.classifier(features * attn) backbone models.resnet50(pretrainedTrue) backbone.fc nn.Identity() # 去掉原分类头 model CustomModel(backbone, 10) # 加载时只取backbone相关的key pretrained_dict models.resnet50(pretrainedTrue).state_dict() model_dict model.state_dict() # 把预训练权重的key加上backbone.前缀来匹配 backbone_dict {fbackbone.{k}: v for k, v in pretrained_dict.items() if not k.startswith(fc.)} model_dict.update(backbone_dict) model.load_state_dict(model_dict)这个前缀映射的操作很关键。因为你的自定义模型把backbone包了一层state_dict的key会变成backbone.conv1.weight这种形式而原始预训练权重是conv1.weight。不做前缀处理一个key都对不上。3.3 场景三层数不一致只加载前N层有些时候你拿到的预训练模型比你需要的深比如用ResNet101的权重去初始化一个ResNet50。这时候只能加载前面能对上的层。pretrained_dict models.resnet101(pretrainedTrue).state_dict() model models.resnet50(pretrainedFalse) model_dict model.state_dict() # 逐层匹配只保留key和shape都一致的 matched_dict {} for k, v in pretrained_dict.items(): if k in model_dict and v.shape model_dict[k].shape: matched_dict[k] v print(f匹配成功: {k}) print(f总共匹配 {len(matched_dict)} 层模型共 {len(model_dict)} 层) model_dict.update(matched_dict) model.load_state_dict(model_dict)这种跨深度加载能匹配上的主要是前面的卷积层和BN层后面的layer3、layer4因为通道数和层数不同基本对不上。实际效果是浅层特征复用深层重新训练。对于数据量不大的任务这种策略往往比完全从头训练收敛快很多。3.4 场景四多GPU训练权重转单GPU用DataParallel或DDP训练保存的权重key前面会多一个module.前缀。加载到单GPU模型时全部对不上。# 多GPU保存的权重 pretrained_dict torch.load(multi_gpu_model.pth) # 去掉module.前缀 new_dict {} for k, v in pretrained_dict.items(): if k.startswith(module.): new_dict[k[7:]] v else: new_dict[k] v model.load_state_dict(new_dict)反过来单GPU权重加载到多GPU模型需要加前缀逻辑类似。这个坑几乎每个人都会踩一次记住module.这个前缀就行。4. 加载完了不等于对了验证与排查的完整链路4.1 加载后必须做的三项检查很多人加载完看到没报错就以为万事大吉结果训练时loss不降反升。我现在的习惯是加载后必做三个检查。第一打印missing_keys和unexpected_keys确认跳过的层符合预期。第二随机抽几个关键层的参数对比加载前后的数值是否真的变了。第三用同一批数据跑一次前向传播看输出是否正常不是NaN量级合理。# 检查一key诊断 result model.load_state_dict(filtered_dict, strictFalse) print(缺失:, result.missing_keys) print(多余:, result.unexpected_keys) # 检查二参数数值对比 for name, param in model.named_parameters(): if name in filtered_dict: loaded filtered_dict[name] if torch.equal(param.data, loaded): print(f{name} 加载成功) else: print(f{name} 数值不一致) break # 检查三前向传播 model.eval() with torch.no_grad(): dummy torch.randn(2, 3, 224, 224) out model(dummy) print(输出形状:, out.shape) print(输出范围:, out.min().item(), out.max().item()) print(是否有NaN:, torch.isnan(out).any().item())4.2 一个真实的排查案例之前有个项目我用EfficientNet-B0做骨干加载官方权重后训练前几个epoch loss正常下降到第5个epoch突然爆炸。排查了半天最后发现问题出在BN层的num_batches_tracked这个buffer上。EfficientNet的预训练权重里BN层的num_batches_tracked是一个很大的数因为在大数据集上训了很久而我的模型初始化时它是0。加载时这个buffer也被加载了导致BN层在计算时使用了预训练阶段的统计量和我的数据分布不匹配。前几个epoch因为momentum更新还没覆盖掉到后面就发散了。解决方案是在加载时过滤掉所有num_batches_trackedfiltered_dict {k: v for k, v in pretrained_dict.items() if num_batches_tracked not in k}这个坑很隐蔽因为num_batches_tracked不是可学习参数打印模型参数时看不到它但它在state_dict里确实存在。如果你用strictFalse它会被静默加载然后埋下隐患。4.3 学习率策略要跟着调整部分加载后不同层的新鲜程度是不一样的。加载了预训练权重的层已经学到了通用特征而随机初始化的层还是白纸。如果对所有层用同一个学习率要么预训练层被破坏要么新层学得太慢。我的做法是分组设置学习率预训练部分用小学习率比如1e-4新初始化的部分用大学习率比如1e-3。在PyTorch里通过参数组实现pretrained_params [] new_params [] for name, param in model.named_parameters(): if name in filtered_dict: pretrained_params.append(param) else: new_params.append(param) optimizer torch.optim.Adam([ {params: pretrained_params, lr: 1e-4}, {params: new_params, lr: 1e-3} ])如果数据量特别小还可以考虑先冻结预训练层只训新层几个epoch再解冻一起微调。这个策略在少样本场景下效果很明显。5. 那些年我踩过的加载坑与应对技巧5.1 坑一key名字差一个字母全部对不上有次用某个开源实现的骨干网络它的模块命名是features.0.weight而我用的官方实现是conv1.weight。结构完全一样就是命名不同导致一个key都加载不上。当时没打印missing_keys训练了半天发现效果和随机初始化一样才回头排查。应对方法加载前先打印两边state_dict的前几个key对比一下。如果发现命名体系不同写一个映射字典手动转换。# 查看两边的key print(预训练权重keys:, list(pretrained_dict.keys())[:5]) print(模型keys:, list(model.state_dict().keys())[:5]) # 如果命名不同建立映射 name_mapping { features.0.weight: conv1.weight, features.0.bias: conv1.bias, # ... } remapped_dict {name_mapping.get(k, k): v for k, v in pretrained_dict.items()}5.2 坑二DataParallel保存的权重加载到CPU模型在GPU上用DataParallel训练保存的权重加载到CPU上推理时除了module.前缀问题还会遇到设备不匹配。torch.load默认会把权重加载到保存时的设备上如果保存时在GPU而你现在只有CPU会报错。# 加载时指定map_location pretrained_dict torch.load(model.pth, map_locationcpu) # 或者加载后统一处理 pretrained_dict torch.load(model.pth, map_locationlambda storage, loc: storage)map_location这个参数建议养成习惯每次都写不管当前有没有GPU。这样代码在不同机器上迁移时不会出问题。5.3 坑三加载了权重但忘了切换train/eval模式这个坑和加载本身无关但经常一起出现。加载预训练权重后如果模型里有Dropout和BatchNorm忘记调用model.train()或model.eval()会导致行为异常。特别是做特征提取时必须model.eval()并且torch.no_grad()否则BN层会更新统计量Dropout会随机丢弃提取的特征不稳定。# 特征提取模式 model.eval() with torch.no_grad(): features model(input_tensor) # 微调模式 model.train() output model(input_tensor)5.4 坑四部分加载后直接保存下次加载又出问题有时候你部分加载后想保存这个半初始化的模型方便下次直接用。但如果保存的是整个模型对象torch.save(model)下次加载时如果代码里的类定义变了会反序列化失败。正确做法是只保存state_dict。# 推荐只保存state_dict torch.save(model.state_dict(), partially_loaded.pth) # 加载时先建模型再加载 model CustomModel(...) model.load_state_dict(torch.load(partially_loaded.pth))这样模型结构和权重分离代码重构时不会因为类定义变化导致加载失败。5.5 一个实用的加载工具函数踩了这么多坑之后我封装了一个通用的部分加载函数每次直接用省得重复写过滤逻辑def load_pretrained_partial(model, pretrained_path, skip_keysNone, verboseTrue, map_locationcpu): 部分加载预训练权重 skip_keys: 需要跳过的key前缀列表如[fc, classifier] pretrained_dict torch.load(pretrained_path, map_locationmap_location) model_dict model.state_dict() # 去掉module.前缀 new_dict {} for k, v in pretrained_dict.items(): if k.startswith(module.): k k[7:] new_dict[k] v pretrained_dict new_dict # 过滤skip_keys if skip_keys: pretrained_dict {k: v for k, v in pretrained_dict.items() if not any(k.startswith(sk) for sk in skip_keys)} # 过滤num_batches_tracked pretrained_dict {k: v for k, v in pretrained_dict.items() if num_batches_tracked not in k} # 只保留key和shape都匹配的 matched_dict {} for k, v in pretrained_dict.items(): if k in model_dict and v.shape model_dict[k].shape: matched_dict[k] v model_dict.update(matched_dict) result model.load_state_dict(model_dict, strictFalse) if verbose: print(f成功加载 {len(matched_dict)}/{len(model_dict)} 层) if result.missing_keys: print(f缺失的key: {result.missing_keys}) return model, result这个函数覆盖了前缀处理、skip过滤、形状匹配、buffer过滤几个常见需求。实际用的时候根据情况传参就行。6. 不同框架版本下的行为差异PyTorch不同版本在load_state_dict的行为上有一些细微差别不注意的话升级版本后可能出问题。在1.0到1.4版本之间strictFalse时如果key存在但shape不匹配会直接报错而不是跳过。从1.5开始shape不匹配的key会被归入unexpected_keys并跳过。所以如果你在旧版本上写的代码依赖了shape不匹配会报错这个行为来做判断升级后逻辑就失效了。另外assign参数是1.13才加入的。默认assignFalse时加载是原地复制数值assignTrue时会用预训练权重的张量替换模型原有的张量。区别在于assignTrue会保留预训练权重的设备、dtype等属性。如果你从GPU权重加载到CPU模型用assignTrue可能会把模型参数变成CUDA张量导致后续CPU计算报错。# 安全做法明确指定map_location不用assign model.load_state_dict(torch.load(path, map_locationcpu), strictFalse)还有一个容易忽略的点torch.load在2.6版本之后默认weights_onlyTrue这意味着它只能加载纯张量字典不能加载包含自定义类的完整模型。如果你保存的是整个模型对象加载时会报错。解决方案是保存时就用state_dict或者加载时显式传weights_onlyFalse但这样有安全风险只加载可信来源的文件。7. 从部分加载延伸到迁移学习的策略选择部分加载只是手段最终目的是让模型在新任务上表现好。加载完之后怎么训其实比怎么加载更影响结果。如果你的新任务数据和预训练数据分布接近比如都是自然图像那加载大部分层、只换分类头、用小学习率微调就够了。如果分布差异大比如自然图像到医学影像浅层特征可能也不适用这时候可以考虑只加载前几层后面全部重训或者干脆从头训练但用预训练权重做初始化。数据量也是个关键因素。数据少的时候冻结预训练层只训新层能防止过拟合数据多的时候全部解冻一起微调效果更好。我一般的做法是先用小学习率全部微调跑几个epoch看验证集表现如果过拟合就冻结浅层如果欠拟合就加大学习率或解冻更多层。还有一个技巧是分层设置学习率衰减浅层学习率最小深层逐渐增大。因为浅层学到的是通用特征边缘、纹理不需要大改深层学到的是任务相关特征需要更多调整。# 分层学习率示例 params [ {params: model.conv1.parameters(), lr: 1e-5}, {params: model.layer1.parameters(), lr: 1e-5}, {params: model.layer2.parameters(), lr: 5e-5}, {params: model.layer3.parameters(), lr: 1e-4}, {params: model.layer4.parameters(), lr: 5e-4}, {params: model.fc.parameters(), lr: 1e-3}, ] optimizer torch.optim.Adam(params)这套组合拳打下来大部分迁移学习场景都能覆盖。核心思路就是加载时精细控制哪些层复用训练时区别对待不同层的更新幅度。我在实际项目里最深的体会是部分加载这件事写代码只占两成时间剩下八成都在验证和调试。每次加载完多花五分钟做检查能省下后面几小时的排查。尤其是missing_keys和unexpected_keys这两个列表一定要看不要嫌麻烦。