ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

解决size mismatch:PyTorch加载预训练模型前如何正确替换全连接层

解决size mismatch:PyTorch加载预训练模型前如何正确替换全连接层 在深度学习里混久了你会发现报错这事也分三六九等。有些错属于低级语法问题瞄一眼就能解决有些错属于环境配置地狱没个半天出不来还有一种错属于“你明明每一步都写了但就是哪里不对”的玄学。我这些年带过不少实习生见过最多的一个报错毫无悬念是size mismatch for fc.weight几乎所有第一次接触预训练模型迁移学习的人都会在这个地方栽一次。这个报错背后其实就藏着一个特别容易被忽略的动作——加载预训练模型前你根本没想起来要处理全连接层。今天我不讲那些大而全的迁移学习理论就专门把这个“关掉全连接层”的事从头到尾掰扯清楚为什么必须关、怎么关、关了之后怎么办以及不同模型结构下有哪些坑等着你。1. 先搞清楚预训练模型和全连接层到底在扮演什么角色1.1 预训练模型一份已经帮你“看”过百万张图的简历先说预训练模型。以计算机视觉里最常用的 ResNet-50 为例它起先在 ImageNet 这个包含 128 万张图片、1000 个类别的大数据集上完成了训练。ImageNet 里的类别五花八门从“金毛犬”“阿富汗猎犬”这种细分类到“蘑菇”“降落伞”这种完全不搭边的类别都有。训练完成后这个模型的前面绝大部分层卷积层、BN 层、残差块就学会了一件非常重要的事情把一张原始图片逐步抽象成一组有意义的特征向量。比如底层卷积可能学会识别边缘、纹理、颜色块中层学会识别眼睛、轮子、叶片这些部件高层学会把这些部件组合成更抽象的语义。很多人喜欢把预训练模型比喻成“预培训过的员工”我觉得不够贴切。我更愿意把它比喻成一份简历——这份简历不是把“能力”写下来的文字而是一整套已经调好的参数是一双被 128 万张图喂过、已经知道怎么“看”世界的眼睛。你拿到它的那一刻人家已经替你完成了从零训练最烧钱的那部分工作。1.2 全连接层一张写满旧答案的答题卡好现在说说全连接层也就是 Fully Connected Layer俗称 FC 层或 Linear 层。在卷积神经网络里FC 层几乎总是出现在网络的最后。前面卷积部分不断把图像“压扁”成特征图到最后一层时这些特征图会被展平Flatten成一维向量然后全连接层拿着这个向量去做最终裁决。全连接层做的事情本质就是一次矩阵乘法加偏置output weight * input bias。这个weight的形状是[输入特征维度, 输出类别数]。在 ImageNet 预训练的 ResNet-50 上这个权重就是[2048, 1000]——因为特征维度是 2048而 ImageNet 要分出 1000 个类别。你可以把 FC 层理解成一张答题卡每个类别对应一个“答案模板”。模型把前面提取的特征向量跟这张答题卡逐一比对算出一个相似度分数分数最高那个类别就是模型的判断结果。1.3 为什么一加载就报错维度就是最现实的坑明白了上面这些你大概能猜到报错的原因了。假设你现在要做的是 10 分类的花卉识别你拿到一个 ImageNet 上训练好的 ResNet-50直接写model torchvision.models.resnet50(weightstorchvision.models.ResNet50_Weights.IMAGENET1K_V1) ckpt torch.load(resnet50_weights.pth) model.load_state_dict(ckpt)然后load_state_dict内部就会开始做“找对象”式的配对它会拿着权重文件里的每一个 key比如fc.weight、fc.bias去模型结构里找同名 key。找到之后还要对比形状。于是你就看到了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]).翻译成人话就是权重文件里的 FC 层是个 1000 类的答题卡你模型里现在的 FC 层是个 10 类的答题卡这俩对不上。这就是“预训练模型 全连接层”组合在一起时最直接、最普遍的问题。2. 为什么必须“关掉”全连接层不只是维度问题很多人觉得那我只要把新任务的类别数改成和 ImageNet 一样都是 1000 类是不是就万事大吉了不是的。这里面的坑比维度不匹配深得多。2.1 语义空间错位1000类学到的边界对N类任务毫无意义深度学习模型的最后一个全连接层学到的不是一个通用的“分类能力”而是一套非常具体的类间边界。它是在 1000 个固定类别之间画出来的决策边界。比如它在“金毛犬”和“拉布拉多犬”之间画了一条线在“蘑菇”和“降落伞”之间画了另一条线。你的新任务如果是“识别 1000 种不同的蘑菇”那原 FC 层的语义空间和你的任务还算沾点边但如果你做的是“猫狗二分类”原 FC 层里根本没有“猫”和“狗”这两个类别的概念——不是没有这个词而是它的参数里存的是 ImageNet 里那 1000 个具体类别的映射权重。哪怕你的新任务恰好也是 1000 类只要这 1000 类和 ImageNet 的不一样原 FC 层对你来说就是一张“答非所问”的废卡。它那些精心调过的权重是对着别人的题目优化出来的拿到你的题目上非但没用反而会带偏梯度。2.2 特征层级的规律越靠后的层越“任务专用”这里我得提一个深度学习中很重要的观察结论在卷积神经网络中越靠近输入的层学到的特征越通用越靠近输出的层学到的特征越专用于当前任务。底层卷积权重基本就是一堆边缘检测器、纹理过滤器换个数据集照样能用。但高层特征和最后的 FC 层已经跟原训练任务的类别标签深度绑定了。这就好比一个人学会了“感知颜色、形状、明暗”这种通用视觉能力然后考了一场 1000 道题的试考完之后脑子里的答案全都是那 1000 道题的。你让这个“脑子”去参加你的新考试正确的做法肯定不是连那 1000 道题的答案一起带过去而是只保留他“看东西”的能力重新教他新答案。这就是“关掉全连接层”的本质只继承通用特征提取部分丢掉任务专用分类头。2.3 留着旧全连接层还有哪些隐性风险除了语义错位留着旧的 FC 层还会有几个实际工程上的麻烦过拟合风险更高。很多经典网络比如 VGG的全连接层占了整个模型参数的绝大部分这些参数如果保留并使用不合适的学习率更新很容易在小型数据集上快速过拟合。梯度传导混乱。老 FC 层在初始化状态下对你的新任务没有任何有意义的梯度方向前几个 epoch 会输出一堆随机概率如果 BatchNorm 等层还在更新可能会导致训练早期剧烈震荡。输出维度上下不匹配。如果你的新任务类别数大于 1000或者不是 1000FC 层权重根本就处理不了。还有些特殊情况比如你的输出不是做分类而是做回归、检测框回归、Embedding那 FC 层输出维度完全需要另算留着旧的就是在添乱。3. 实操PyTorch 里怎么正确加载并替换全连接层说完了道理直接上干货。下面这部分是你在 PyTorch 里“关掉”全连接层的完整套路我按不同情况分开说。3.1 动手之前先确认三件事别急着敲代码。先确认以下三点你的任务是几分类num_classes是什么如果是回归任务那输出维度是多少你用的预训练模型是哪个家族的ResNet 的最后一层叫model.fcVGG 和 MobileNet 叫model.classifierViT 叫model.head。名字都不同替换的时候别找错门。前面的特征输出维度是多少比如 ResNet-50 是 2048ResNet-18 是 512VGG16 的classifier[6]输入是 4096。不知道的话可以直接通过model.fc.in_features之类的方式动态获取别硬编码。3.2 三种替换全连接层的写法写法一ResNet 家族直接换fcimport torch import torch.nn as nn import torchvision.models as models num_classes 10 # 加载预训练模型 model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # 关键技巧用 in_features 动态获取原fc层的输入维度避免硬编码 in_features model.fc.in_features # 替换成自己的分类头 model.fc nn.Linear(in_features, num_classes)就这么简单。model.fc.in_features这个技巧网上很多教程都不提但它非常实用——你换不同深度的 ResNet18、34、50、101时代码不用改一行。写法二VGG 家族的classifierVGG 的结构跟 ResNet 不一样它的分类头是三个全连接层串在一起存在model.classifier这个Sequential里面model models.vgg16(weightsmodels.VGG16_Weights.IMAGENET1K_V1) # 把 classifier 的最后一层替换掉 num_features model.classifier[6].in_features model.classifier[6] nn.Linear(num_features, num_classes)注意VGG 的classifier前两个 Linear 是 4096 维的隐藏层一般建议保留它们也可以作为通用特征提取器的一部分只替换最后输出层。写法三更自由的做法自己定义一个全新的“头”有时候你不想只换一层。比如想在分类头里加一层 Dropout或者想先接一个全局平均池化再分类可以自己定义一个nn.Sequential或子网络直接替换class CustomHead(nn.Module): def __init__(self, in_features, num_classes): super().__init__() self.global_pool nn.AdaptiveAvgPool2d((1, 1)) self.dropout nn.Dropout(0.2) self.fc nn.Linear(in_features, num_classes) def forward(self, x): x self.global_pool(x) x torch.flatten(x, 1) x self.dropout(x) return self.fc(x) model.fc CustomHead(in_featuresmodel.fc.in_features, num_classes10)3.3 load_state_dict 的坑和正确打开方式替换完 FC 层之后你直接model.load_state_dict(ckpt)还是会报错因为权重文件里的fc.weight和模型里的fc.weight形状不同。这时候很多人的第一反应是用strictFalsemodel.load_state_dict(ckpt, strictFalse)这里我要提醒你strictFalse确实可以跳过不匹配的 key但它没那么智能。如果你新任务的类别数恰好也是 1000这个参数会直接误把 ImageNet 的 FC 层权重加载进来而你可能压根不想要它。更稳妥的做法是先过滤掉 FC 层相关的 key再加载ckpt torch.load(resnet50_weights.pth, map_locationcpu) # 过滤掉 fc 层的权重 filtered_ckpt {k: v for k, v in ckpt.items() if not k.startswith(fc.)} # 先加载过滤后的预训练权重 model.load_state_dict(filtered_ckpt, strictFalse) # 再替换 FC 层 model.fc nn.Linear(in_features, num_classes)顺序上我建议先load_state_dict再替换fc这样结构的变化不影响加载。如果你的模型里还有其他自定义层也可以用同样的方式过滤调对应 key。3.4 替换之后的训练策略冻结还是微调把旧的 FC 层“关掉”之后新 FC 层是一个随机初始化的空壳接下来怎么训练直接决定最终效果。常见做法有两种。方案 A只训练新分类头迁移学习适用于数据集比较小、跟 ImageNet 分布差别不大的场景。冻结 backbone只更新自己的 FC 层# 冻结所有层 for param in model.parameters(): param.requires_grad False # 解冻新的 FC 层 model.fc nn.Linear(in_features, num_classes) for param in model.fc.parameters(): param.requires_grad True # 优化器只需要传入 fc 层的参数 optimizer torch.optim.Adam(model.fc.parameters(), lr1e-3)方案 B全量微调但分层设置学习率适用于数据量较大、或者域差异明显的场景。backbone 用一个小学习率去轻微调整新 FC 层用高一点的学习率去快速学习optimizer torch.optim.SGD([ {params: [p for n, p in model.named_parameters() if fc not in n], lr: 1e-5}, {params: model.fc.parameters(), lr: 1e-3}, ], momentum0.9, weight_decay1e-4)我个人的习惯是在替换 FC 层后先用方案 A 跑几个 epoch让新分类头先找到一个大致像样的决策区域然后再解冻 backbone 做微调。这个思路类似于先教“新员工”认清试卷格式再让他去理解出题风格通常比从头就全量微调要稳定很多。注意新 FC 层是随机初始化的它在前几个 epoch 的输出接近于均匀分布的随机概率。这时候如果 backbone 也在同步大范围更新整个模型可能被带偏。所以不管选哪个方案新 FC 层的学习率都不要太小backbone 的学习率不要太大这是一个经验值。4. 延伸思考全连接层的替代品与 NLP 里的类似操作4.1 ResNet 的 GAP 思路把 FC 层“拍扁”来用既然全连接层这么麻烦能不能不用它可以。ResNet 论文本身其实已经做了一个很重要的设计决策在网络的最后用一个全局平均池化Global Average PoolingGAP替代了 VGG 那种多层全连接堆叠的分类头。ResNet 的结构是这样的最后一个卷积层输出的特征图尺寸是[2048, 7, 7]接下来不是直接展平接一个 2048 → 1000 的 FC 层而是先做 GAP把每个通道的 7×7 特征图平均成一个值得到[2048]的向量然后才接一个 FC 层。GAP 的好处非常明显它把“空间位置”的信息直接给抹平了对物体的平移、形变更鲁棒同时因为 FC 层的参数量稀释到只剩一层线性层过拟合风险大大降低。你替换分类头的时候如果是在 ResNet 上操作其实已经默认享受了这个好处。我遇到过一些朋友在替换 FC 层时还会纠结要不要自己再加一个Flatten。在 ResNet 上其实不需要因为前后view逻辑 PyTorch 已经帮你处理好了你只管替换fc就行。但如果你在自定义网络里用了 GAP别忘了在进入 FC 前把多维特征展平成一维。4.2 径向核函数做分类头一个“看着很美落地有坑”的方案再说回热搜里的“径向核函数代替全连接层”。径向基函数RBF是一种经典的核函数常见形式是高斯函数exp(-||x - c||² / (2σ²))其中c是中心点σ是宽度参数。RBF 分类器的思路是把输入样本映射到以某些中心点为基准的相似度空间中然后用这些相似度分数做分类。理论上拿 RBF 分类器当网络最后的“头”替换全连接层有一个天然优势非线性表达能力更强能刻画更复杂的决策边界而且对输入分布的局部结构更敏感。但为什么深度学习时代很少看到有人这么干因为难点全在工程上中心点c怎么选如果固定为训练集的聚类中心那新样本怎么处理如果学出来初始化和训练稳定性都很难保证。宽度σ怎么确定这个参数直接控制核函数的“作用半径”它对最终分类边界影响极其敏感但用梯度下降优化它非常容易陷入局部最优。批量训练的冲突。RBF 网络天然适合基于距离的优化目标而现代深度模型普遍配合交叉熵损失做端到端训练两者的梯度尺度差异很大很难调和。所以在实际项目中RBF 做分类头更多停留在研究和特定场景的尝试。我的建议是如果你不是在做实验探索还是用标准的 FC 分类头更省心。可以把 RBF 当作知识去了解但别轻易在生产环境里拿来当主力。4.3 NLP 里也要“关掉”全连接层以 RoBERTa 中文预训练模型为例聊完视觉我们把视野拉到自然语言处理。如果你用过 Hugging Face 加载中文 RoBERTa 预训练模型比如hfl/chinese-roberta-wwm-ext会发现同样的坑在 NLP 里一个不少。BERT、RoBERTa 这类模型的预训练任务通常是“掩码语言模型”MLM和“下一句预测”。它们顶层也有一个全连接层但那是为预训练任务服务的输出的是词表大小比如 21128的概率分布。你如果要做文本情感二分类这个顶层输出维度压根对不上。Hugging Face 的做法是提供一个AutoModelForSequenceClassification它会自动帮你把预训练模型的顶层分类头替换成一个适配num_labels的新 Linear 层from transformers import AutoModelForSequenceClassification model AutoModelForSequenceClassification.from_pretrained( hfl/chinese-roberta-wwm-ext, num_labels2 )这里的逻辑和视觉那边完全一致预训练权重只加载到底层编码器顶部分类头重新初始化。唯一的区别是视觉模型里 FC 层的名字叫fc或classifierNLP 模型里叫classifier或者pooler但本质工作都是“关掉旧分类头、换上新高”。如果你自己手动修改 RoBERTa 的分类头比如把两分类改成六分类也别忘了维度要匹配否则同样会蹦出 size mismatch 的报错。这再次说明“加载预训练模型前处理掉顶层”是跨模态、跨领域的通用原则。5. 常见问题排查与操作心得5.1 加载预训练模型时最常踩的 5 个坑我把这些年看到过、也亲自踩过的问题整理成一张速查表方便你对症下药报错或现象根本原因解决办法size mismatch for fc.weight新旧 FC 层输出维度不一样替换 FC 层后再加载权重或过滤掉 fc 相关 keyMissing key(s) in state_dict: fc.weight权重文件里有 fc但模型里已经被你删了或换名了用 strictFalse或者把权重文件里的 fc key 过滤掉Unexpected key(s) in state_dict: fc.weight权重文件里的 fc key 在模型里找不到确认模型结构是不是匹配的版本必要时过滤 key加载成功但前向传播报维度错误输入特征维度和 FC 层权重不一致检查 backbone 的输出维度用in_features动态创建 FC 层训练起来 loss 不下降或剧烈震荡旧 FC 层权重复用了或新 FC 层学习率太大/太小确认 FC 层被替换成随机初始化的新层按 backbone 和 FC 分层设置学习率5.2 几条实操心得最后分享几个平时不会写在系统教程里、但实际帮了我大忙的小习惯。第一加载完成后先用一个假的输入跑一次前向。不要急着训先用一个随机 Tensor 试一下dummy_input torch.randn(2, 3, 224, 224) output model(dummy_input) print(output.shape) # 期望输出 [2, num_classes]这一步十秒钟就能做完能帮你提前发现绝大多数学不到的“隐性维度不匹配”问题。第二保存模型时把 backbone 和 FC 层分开留档。我在做项目时习惯保存两种权重一是完整模型backbone 自己的 FC 层二是仅包含 backbone 部分的权重key 过滤后再存。这样下次换任务、换分类头时我可以直接复用 backbone 的权重不用再重新加载一个完整模型再去过滤。# 保存 backbone-only 权重的示例 backbone_state {k: v for k, v in model.state_dict().items() if fc not in k} torch.save(backbone_state, backbone_only.pth)第三别在 FC 层后面手动加 Softmax。绝大多数 PyTorch 训练代码用的是nn.CrossEntropyLoss这个损失函数内部已经包含了 Softmax NLLLoss。你要是手痒在输出层前加了一个nn.Softmax()也不一定会报错但训练出来的结果会变得奇奇怪怪而且你还不容易察觉原因。二分类场景如果用BCEWithLogitsLoss也一样模型输出 logits 就行激活函数交给损失函数处理。第四新 FC 层的初始化其实值得花一点心思。很多人直接nn.Linear用默认初始化在大多数场景下也没问题。但如果你发现训练初期 loss 掉得太慢可以试试用更大范围的 Xavier 初始化或者调高新 FC 层的学习率。最直接的做法就是给 FC 层单独设一个比 backbone 大 10 到 50 倍的学习率效果往往立竿见影。第五注意不同版本的 torchvision API。老版本里写的是model models.resnet50(pretrainedTrue)新版本已经改成weights...的写法新写法还带警告信息。如果你看到别的博客代码跑不通大概率是版本不同导致 API 差异别急着怀疑自己的逻辑。说到最后我想起自己第一次遇到这个报错时的情形。那会儿我刚接触迁移学习在一个十类的花卉数据集上用 ResNet-18死活就是加载不进去权重一运行就报 size mismatch。我在网上翻了差不多两个小时最后看到一条回复写着“你把模型最后一层换掉再加载”我改了代码三秒钟跑通。当时我整个人先是愣住然后是又好气又好笑——折腾了两小时原来就缺了这一行替换代码。从那以后我形成了一个肌肉记忆拿到任何预训练模型第一件事不是load_state_dict而是先看一眼最后一层是什么想清楚要不要换。这个习惯帮我省下了无数个 Debug 的夜晚。如果你现在也卡在这个报错上希望这篇东西能让你少走一点弯路至少下次看到 fc.weight 的时候你能微微一笑知道老朋友又来打招呼了。
返回列表