直接上手:理解 PyTorch 修改预训练模型的三种核心场景
很多人拿到预训练模型时都会遇到这种尴尬:模型在 ImageNet 上效果很好,但自己的分类任务类别数不对,或者想改一下网络中间层的结构,让模型更适应自己的数据分布。因为完整训练一个模型成本太高,特别是在数据量不足的情况下,大家普遍的做法是“站在巨人的肩膀上”——加载预训练权重,再对网络结构做针对性的手术。
这篇文章只讲一件事:在 PyTorch 中如何修改预训练模型的结构,具体拆成三种场景——全连接层修改、单个卷积层替换、多个卷积层连续改造。我会先说明每种操作的适用场景和底层逻辑,再给可直接运行的代码,最后聊一聊我在实际调试中踩过的坑。内容面向已经会用 PyTorch 加载模型、但不太敢动模型内部结构的开发者,读完你应该能独立完成权重迁移和结构替换。
为什么说这个问题值得单独写一篇?因为 PyTorch 加载预训练模型的惯用写法是model.load_state_dict(state_dict),但这句话有个前提:模型结构和保存权重时的结构必须完全一致。一旦你改了全连接层,输出维度变了,直接加载必然报size mismatch错误。卷积层的修改更隐蔽——如果有人只看错误信息,会以为是加载流程出了问题,实际却是网络结构被改动后参数形状对不上而已。
所以我想先把核心逻辑讲透:模型结构决定了参数字典的 key 和 value 形状,改动结构等于改动参数字典的“映射关系”,接下来所有的操作都是在修正这种错位。如果你能理解这句话,后面的代码就是在具体落实这个思路。
1. 修改预训练模型的三种场景及其适用逻辑
1.1 全连接层修改:最常规的迁移学习场景
全连接层修改是所有人首先会遇到的操作。无论是在 PyTorch 官方仓库里下载的 ResNet、VGG,还是用别人开源的检测模型,最后几层基本都是nn.Linear做分类输出。预训练模型的 1000 类输出对应的是一个形状为[1000]的权重矩阵,而你的任务可能只有 10 类或者 2 类,这时就必须把最后一层换掉。
我见过很多初学者在这里犯同一个错误:直接用model.fc = nn.Linear(2048, 10)改了结构,然后运行model.load_state_dict(torch.load('weight.pth')),报错了就怀疑权重文件损坏。实际上 PyTorch 加载权重的逻辑非常简单——它不看层的名字,只看state_dict里的 key 是否一样、value 形状是否一致。当你把fc的输出从 1000 改成 10 时,新层的 key 同样是fc.weight,但形状变成了[10, 2048],和预训练权重里的[1000, 2048]不匹配,于是报size mismatch for fc.weight。
全连接层修改的适用场景很广:
- 分类任务迁移学习:换掉最后的 1000 类分类头,改成自己的类别数。
- 主干网络特征提取:取掉全连接层,把输出变成特征向量,用于检索或者对比学习。
- 多任务学习:把原本的一个分类头改成多个分支,每个分支负责不同任务。
1.2 单个卷积层修改:局部结构优化的起点
第二个场景是修改单个卷积层。有人会问,卷积层又不是最后输出层,改了它有什么用?其实用途很明确:比如你把预训练模型某个 stage 的 3x3 卷积换成了 5x5 卷积,或者把普通卷积换成深度可分离卷积,再或者改变某一层的通道数使其更适配自己的数据规模。
卷积层的修改之所以比全连接层难,是因为它牵一发动全身。假设网络内部某个卷积层原本是Conv2d(256, 512, kernel_size=3),你把输出通道改成 256,那么这一层的参数 shape 变了;更重要的是,下一层卷积层的输入通道也必须跟着改成 256,否则模型在 forward 计算时直接报维度错误。这种连锁反应决定了单个卷积层的修改通常不是“改一层”,而是“改一层并适配相邻层”。
实际项目中我遇到最多的情况是:某个任务输入图片分辨率很高,希望网络前几层的感受野更大一些;或者某个轻量化部署场景,希望把标准卷积替换成分组卷积以减少参数量。这些都属于单个卷积层的修改范畴。
但有一点必须提醒:修改预训练模型的中间层,意味着这一层以及它之后的层的预训练权重基本都失效了。你通常需要冻结前面的层、只微调改动层及后续层,才能让训练稳定下来。
1.3 多个卷积层修改:深层次网络重构
多个卷积层的修改是最复杂也最能体现“手术水平”的一块。常见于:
- 把 ResNet 的
BasicBlock整体替换成Bottleneck结构。 - 在多个 stage 中同时调整通道数,构建更宽或更窄的变体网络。
- 在原有卷积之后插入注意力模块(如 SE Block、CBAM),需要连续改写多个层的前后衔接。
- 去掉网络中冗余的阶段,或者增加新的阶段进行多尺度特征融合。
多个卷积层连续修改时,你的注意力不应该只放在“改哪几层”上,还要考虑整个前向传播路径的数据流——每一层的输入输出 shape 是否首尾相接。我的建议是先把原模型完整打印出来,用print(model)查看每一层的名字和参数,再在纸上画出你要修改的部分涉及的数据流,最后才开始动手。
注意:多个卷积层修改时,强烈建议先把网络的
state_dict完整备份出来再动手。一旦出错,你可以随时恢复到修改前的状态,避免反复下载权重文件。
2. 核心前置知识:参数字典、层名与权重加载机制
2.1 理解 state_dict 的结构
PyTorch 中所有的可学习参数(权重和偏置)都存放在state_dict这个有序字典里。它的 key 是层的完整路径,value 是torch.Tensor。例如 ResNet50 的最后一层全连接层的 key 就是fc.weight和fc.bias;而某个中间卷积层的 key 可能是layer2.0.conv1.weight。
这里的关键是:key 的命名规则完全由模型定义时的模块嵌套决定。如果你在模型里自定义了一个类MyBlock,并且在这个类内部有self.conv = nn.Conv2d(...),那么它的 key 就是网络路径.MyBlock路径.conv.weight。理解了这个规则,你就明白了为什么修改模型后权重的 key 对不上——因为你改了结构,key 的路径也随之改变了。
2.2 load_state_dict 严格模式与宽松模式
load_state_dict默认是严格模式,要求传入的字典和模型当前参数“完全匹配”。只要有一个 key 多出来或者少了一个,都会直接报错。报错信息通常会分成两类:Missing key(s)(模型有但权重里没有)和Unexpected key(s)(权重里有但模型没有)。
在实际修改预训练模型时,严格模式几乎总是会报错的,除非你只改了最后一层的输出维度并且手动删除了对应 key。所以大家普遍的做法是先load_state_dict(state_dict, strict=False)把能匹配上的权重加载进去,然后再检查报错信息,确认缺少的 key 是不是你故意改掉的那些层。
提示:
strict=False不是银弹。它会把所有形状匹配的参数都加载进去,形状不匹配的参数会静默跳过(只给出警告)。如果你改了某一层的输出通道数,相邻层的权重形状也可能不匹配,这些都会被跳过,导致模型部分层是随机初始化。所以在加载后务必打印一遍model.state_dict()和加载结果的对比,确认哪些层用了预训练权重,哪些层是重新初始化的。
还有一个小技巧:当你只想加载部分层时,可以手动构造过滤后的字典。例如只加载主干网络、不加载分类头,可以用一个循环遍历state_dict.items(),根据 key 的前缀筛选后再传入。
3. 实操:三种修改方式的完整代码与说明
3.1 修改全连接层并正确加载权重
我们以 ResNet50 为例。任务是从 1000 类改成 10 类。
import torch import torch.nn as nn import torchvision.models as models # 1. 加载原始模型结构 model = models.resnet50(pretrained=False) # 2. 修改全连接层 num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 10) # 3. 加载预训练权重(不包含分类头) state_dict = torch.load('resnet50_pretrained.pth') # 这种方法假设你的权重文件本身就没有 fc 层 model.load_state_dict(state_dict, strict=False)如果你下载的权重是完整包含 1000 类分类头的,那你需要删掉fc.weight和fc.bias这两个 key 再加载:
state_dict = torch.load('resnet50_pretrained.pth') # 删除分类头对应的 key state_dict.pop('fc.weight', None) state_dict.pop('fc.bias', None) model.load_state_dict(state_dict, strict=False)为什么推荐“删掉再加载”而不是“直接加载后再覆盖”?因为直接加载完整权重会因为 shape 不匹配报错,即便strict=False也只能跳过,同时控制台会刷出好几条警告,影响判断。手动删除 key 更干净,也让后续调试时少一些噪音。
在修改完成后,建议验证一下新模型是否正常:
# 构造一个假输入 dummy_input = torch.randn(1, 3, 224, 224) # 前向传播 output = model(dummy_input) print(output.shape) # 期望输出 (1, 10)如果你发现输出形状不对,多半是修改全连接层时in_features和主干网络实际输出的特征维度不一致。这时候可以打印model.fc.in_features来检查。
3.2 替换单个卷积层并保留其余权重
假设我们要把 ResNet18layer1.0.conv1这个卷积从 3x3 换成 5x5,且保持输入输出通道不变。这种修改不影响前后层的数据流,是最简单的一种。
import torch import torch.nn as nn import torchvision.models as models model = models.resnet18(pretrained=True) # 查看原始卷积层的参数 old_conv = model.layer1[0].conv1 print(old_conv.weight.shape) # 例如 torch.Size([64, 64, 3, 3]) # 构造新的卷积层 new_conv = nn.Conv2d( in_channels=64, out_channels=64, kernel_size=5, stride=old_conv.stride, padding=2, # 需要根据 kernel_size 调整,保证输出尺寸不变 bias=old_conv.bias is not None ) # 替换 model.layer1[0].conv1 = new_conv # 加载权重时,删掉旧的和新的 key 之间不匹配的部分 state_dict = model.state_dict() # 注意:现在 model 已经是修改后的注意,上面的代码最后一步其实没有加载预训练权重,只是保存了修改后的随机初始化模型。真正要做的是先打印出替换前后state_dict中层的名字,然后做对应替换。
这里有个更通用的做法——先加载预训练权重到原模型上,再修改结构,最后把权重“搬运”到新层上:
# 1. 先加载预训练模型(原始结构) model_old = models.resnet18(pretrained=True) state_dict_old = model_old.state_dict() # 2. 修改结构 model_new = models.resnet18(pretrained=False) new_conv = nn.Conv2d(64, 64, kernel_size=5, stride=1, padding=2) model_new.layer1[0].conv1 = new_conv # 3. 把能匹配的权重都加载进去 state_dict_old.pop('layer1.0.conv1.weight', None) state_dict_old.pop('layer1.0.conv1.bias', None) model_new.load_state_dict(state_dict_old, strict=False)这种“先加载后修改再搬运”的方式,比“修改后加载”好在哪?因为修改后的模型已经无法直接加载完整预训练权重了,而你需要在修改前就把预训练参数拿到手,再手动迁移到新层里。这个思路是所有复杂模型修改的地基。
想更进一步,如果你期望保留原 3x3 卷积的权重复制到新的 5x5 卷积中(中心填充),可以做一次张量复制:
with torch.no_grad(): new_conv.weight.zero_() # 5x5 卷积核的中心 3x3 区域,填充原始权重 new_conv.weight[:, :, 1:4, 1:4] = old_conv.weight if old_conv.bias is not None: new_conv.bias.copy_(old_conv.bias)这个技巧在修改感受野时非常实用,因为网络能继承原有的低级特征提取能力,而不必从头学起。
3.3 多个卷积层批量替换与权重迁移
多个卷积层修改时,一个典型的场景是把 ResNet 的BasicBlock换成Bottleneck。这种结构变化会改变 stage 内部的层数、卷积通道数、shortcut 结构,几乎相当于重写整个 stage。
我个人更推荐一种更工程化的做法——定义一个修改后的新模型类,然后利用load_state_dict的 key 匹配特性,把可用权重加载进去。步骤是这样的:
import torch import torch.nn as nn import torchvision.models as models # 定义一个修改后的模型骨架 class ModifiedResNet(nn.Module): def __init__(self, num_classes=10): super().__init__() # 这里可以使用官方 resnet 作为主干,也可以自定义 self.backbone = models.resnet18(pretrained=False) # 修改多个卷积层,例如替换前两个 stage 的 conv1 # 重新定义层时保持命名一致很重要 self.backbone.layer1[0].conv1 = nn.Conv2d(64, 64, 5, padding=2) self.backbone.layer1[1].conv1 = nn.Conv2d(64, 64, 5, padding=2) # ... 替换最后一个全连接层 self.backbone.fc = nn.Linear(512, num_classes) def forward(self, x): return self.backbone(x) model_new = ModifiedResNet(num_classes=10) # 加载预训练权重(原始 ResNet18 的权重) state_dict = torch.load('resnet18.pth', map_location='cpu') # 删掉不匹配的 key for k in list(state_dict.keys()): if k.startswith('layer1.0.conv1.') or k.startswith('layer1.1.conv1.') or k.startswith('fc.'): del state_dict[k] model_new.load_state_dict(state_dict, strict=False)这不是唯一方案,但可以说是最不费脑的方案。核心思想是:把模型的结构修改封装成类,修改后还可以继续细调;权重的迁移逻辑完全用 key 名驱动,不需要手动一个个赋值。
如果修改的层非常多,手工写一个删除 key 的列表不够灵活。我一般这么处理:先打印出model_new.state_dict()的所有 key,再对照原始权重的 key,写一个函数自动求出两个集合的差集。凡是差集涉及的 key,都在加载前删掉。具体代码如下:
def load_partial_state_dict(model, state_dict, verbose=True): model_keys = set(model.state_dict().keys()) state_keys = set(state_dict.keys()) # 差集:模型里有但权重里没有 missing = model_keys - state_keys # 差集:权重里有但模型里没有 unexpected = state_keys - model_keys # 形状不匹配的 key mismatch = [] for k in (model_keys & state_keys): if model.state_dict()[k].shape != state_dict[k].shape: mismatch.append(k) if verbose: print('Missing:', sorted(missing)) print('Unexpected:', sorted(unexpected)) print('Shape mismatch:', sorted(mismatch)) # 剔除所有需要跳过的 key filtered = {k: v for k, v in state_dict.items() if k in model_keys and k not in mismatch} model.load_state_dict(filtered, strict=False)这个函数我几乎在每个迁移学习项目里都会复用。它能让你在修改完模型后一目了然地看到哪些层用了预训练权重,哪些层是随机初始化。
3.4 修改后新层的初始化策略
把所有删掉的层交给“随机初始化”其实不是一个好习惯。PyTorch 默认的卷积初始化是 Kaiming 均匀分布,对于小数据集微调来说通常够用;但如果有机会,给改动层设置一个更好的初始化能明显加速收敛。推荐做法是:从主干的预训练权重中提取形状兼容的部分来初始化,实在提取不了就保持 PyTorch 默认初始化,或者用nn.init.kaiming_normal_手动初始化。
下面是一次多卷积层修改后的初始化例子:
def init_modified_layers(model): for name, module in model.named_modules(): if isinstance(module, nn.Conv2d): nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu') if module.bias is not None: nn.init.constant_(module.bias, 0) elif isinstance(module, nn.BatchNorm2d): nn.init.constant_(module.weight, 1) nn.init.constant_(module.bias, 0)要注意:对已经成功加载预训练权重的层,千万不要再调用初始化函数,否则会把预训练参数覆盖掉。所以在执行初始化之前,建议先用上面的load_partial_state_dict加载完权重,再对missing和mismatch列表中的层做初始化。
4. 实操过程与验证链路:修改模型前后的完整流程
4.1 修改模型前必须做的四件事
我总结了一套固定的“术前检查清单”,每次修改预训练模型都会走一遍:
- 打印原始模型结构:
print(model),找到所有要修改层的精确名字。 - 保存原始
state_dict:torch.save(model.state_dict(), 'origin.pth')。 - 构造假输入做前向传播:记录输出 shape,特别是主干网络最后输出的特征维度。
- 确认修改后的模型能正常前向传播再谈加载权重。
这四条缺一不可。尤其是第四件事——很多人改完模型直接训练,模型连 forward 都跑不通,报错信息却总是指向权重加载,让人误判问题方向。
4.2 权重加载与验证策略
以一个 ResNet50 修改全连接层并替换layer3中两个卷积层为例。完整的实操流程如下:
import torch import torch.nn as nn from torchvision import models # Step 1: 加载原始模型 model = models.resnet50(pretrained=True) # Step 2: 备份旧权重 old_state_dict = model.state_dict() # Step 3: 修改结构 model.fc = nn.Linear(2048, 5) model.layer3[1].conv2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) model.layer3[2].conv1 = nn.Conv2d(512, 1024, kernel_size=1, stride=1, bias=False) # Step 4: 构造假输入跑通 forward x = torch.randn(2, 3, 224, 224) try: out = model(x) print('Forward OK, output shape:', out.shape) except Exception as e: print('Forward failed:', e) # 这里需要根据报错信息检查是哪一层的维度不匹配 raise # Step 5: 加载权重(去掉修改层的 key) new_state_dict = model.state_dict() # 找出所有需要删除的 key removed_keys = ['fc.weight', 'fc.bias', 'layer3.1.conv2.weight', 'layer3.1.conv2.bias', 'layer3.2.conv1.weight'] for k in removed_keys: old_state_dict.pop(k, None) model.load_state_dict(old_state_dict, strict=False) # Step 6: 验证参数是否真的加载成功 for name, param in model.named_parameters(): if 'layer1' in name or 'layer2' in name: assert not torch.all(param == 0), f'{name} is empty' print('Pretrained weights loaded correctly')这段代码里的 Step 4 非常关键。有的人会说“print 模型结构不就能看出维度吗”,但在复杂网络中,人眼检查不如机器跑一次来得可靠。一次成功的 forward 至少能证明你的修改没有破坏数据流。
4.3 训练前的最后检查:哪些层需要冻结
修改完模型后,多数项目会进入微调阶段。这时你要决定哪些层冻结、哪些层参与训练。我的通用经验是:
- 如果数据集很小(几百到几千张),冻结主干的前 80% 层,只训练改动层和最后几层。
- 如果数据集中等(几万张),可以放开所有 BatchNorm 层参与训练,卷积层前 50% 冻结。
- 如果数据集够大,那就全量微调,但改动层需要设置更高的学习率。
一个常见的参数冻结写法:
for name, param in model.named_parameters(): if 'layer1' in name or 'layer2' in name: param.requires_grad = False冻结层的原因在于:预训练模型靠前面的层提取通用特征(边缘、纹理、颜色),这些特征在不同数据集间是共享的。你改动的层主要在后半部分,负责的是语义信息的整合,这部分需要用自己的数据重新学习。
5. 常见报错与排查技巧实录
5.1 经典报错 size mismatch
这个报错几乎每个人都会遇到。报错信息一般长这样:
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]).处理方式非常简单:不是改模型去迎合权重,而是根据你的预期判断哪个是“正确”的。如果你想保留预训练分类头的 1000 类输出能力,那模型就不该改输出层;如果确实想改成 10 类,那就接受这一层权重被随机初始化的事实,从state_dict里删掉对应 key。
5.2 Missing key(s) 和 Unexpected key(s) 同时出现
当打印出如下信息时,说明模型结构严重偏离了权重:
Missing key(s) in state_dict: layer4.0.conv1.weight, layer4.0.conv2.weight Unexpected key(s) in state_dict: layer4.1.conv1.weight, layer4.1.conv2.weight这种一般发生在多个卷积层连续修改后,stage 的编号或路径变了。例如你把layer4里的模块结构调整了,原本的layer4.0在删除或新增模块后可能变成了layer4.2,key 的路径就完全对不上了。
排查思路是:打印出两份完整的 key 列表,逐行对比。我通常会导出成 txt 文件再用 vscode 对比,比在终端里翻效率高得多。
with open('model_keys.txt', 'w') as f: for k in model.state_dict().keys(): f.write(k + '\n') with open('weight_keys.txt', 'w') as f: for k in torch.load('weight.pth').keys(): f.write(k + '\n')5.3 加载权重不报错,但模型效果极差
有一种更隐蔽的情况:load_state_dict(strict=False)没有报错,但你发现训练几个 epoch 后 loss 下降非常慢,甚至不如随机初始化。问题多半出在:关键层权重被静默跳过了,模型大部分卷积层都是随机初始化状态,只有少部分层加载成功。
解决办法是加载后必须主动检查:
# 打印所有参与加载的层数量 loaded = 0 skipped = 0 for k in state_dict: if k in model.state_dict() and model.state_dict()[k].shape == state_dict[k].shape: loaded += 1 else: skipped += 1 print(f'Loaded: {loaded}, Skipped: {skipped}')只要skipped的数量超过你预期修改的层数,就说明还有隐藏的形状不匹配问题,需要回头检查。
5.4 修改卷积层后 BatchNorm 参数被误删
很多人修改卷积层时只关心conv的权重,忘了同一模块里 BatchNorm 层的参数 key 也存在连带关系。例如 ResNet 的BasicBlock里有bn1和bn2,它们和卷积层是焊接在一起的。如果只替换了conv1而没有移除对应的bn1,那bn1的 key 还在,但它的输入通道可能已经变了(因为conv1输出通道变了),前向传播时并不会报错——BatchNorm 是逐通道归一化,输入通道多少它都能算——但训练时梯度回传就全都乱了。
我的建议是:每当你替换一个卷积层时,把同一个 block 内与之相连的 BatchNorm 层也一起换掉或者重新初始化。不要只盯着卷积权重。
完整的一条经验总结:
修改预训练模型的本质,是保证“结构改变后仍能最大化利用原有参数知识”。对全连接层,删掉对应 key 即可;对单个卷积层,可以考虑在原卷积核中心填充复用参数;对多个卷积层,必须从上到下梳理数据流,逐一确认 key 的匹配。
5.5 一个小技巧:如何快速验证新模型和旧模型在相同输入下的输出差异
当你只是修改了某一层的感受野或通道数,但大体结构没变时,可以做一个“输出对比实验”:用同一个随机输入,分别在旧模型(已加载预训练权重)和新模型(部分层用预训练权重)上跑前向传播,把输出特征图的 L2 距离打印出来。
model_old.eval() model_new.eval() x = torch.randn(4, 3, 224, 224) with torch.no_grad(): feat_old = model_old(x) feat_new = model_new(x) diff = (feat_old - feat_new).norm() / feat_old.norm() print(f'Relative difference: {diff.item():.4f}')如果相对差异小于 0.1,说明你的修改没有严重破坏原有特征提取能力,可以放心进入微调。如果差异特别大,说明你改动的层处于非常浅的位置,连低级特征都变了,这时候最好只改网络深层的结构。
我在实际项目里反复验证过这个技巧——它几乎成了我每次修改预训练模型的“安全阀”。
6. 实操经验补充:常见问题速查表
| 问题现象 | 可能原因 | 解决方案 |
|---|---|---|
| 加载权重报 size mismatch | 修改层和权重 key 形状不一致 | 用strict=False定位不匹配 key,再决定删除或保留 |
| 加载报 Missing key(s) | 新模型有层但权重里没有 | 检查是否新增了层,确认后接受随机初始化 |
| 加载报 Unexpected key(s) | 权重里有层但新模型没有 | 确认旧权重是否来自不同结构的模型 |
| 模型能 forward 但 loss 不降 | 关键层权重未成功加载 | 打印 loaded/skipped 数量,检查是否有隐藏不匹配 |
| 修改卷积层后前一输出正常,后一层报错 | 相邻层通道数未同步调整 | 检查所有与该层相连的层,同步修改通道配置 |
| 修改后模型参数量异常膨胀 | kernel_size 过大或通道数设置不当 | 打印每层参数数量,定位异常层 |
还有几个细节值得单独说:
- 修改模型时尽量用
copy.deepcopy保留原始模型实例,方便后续比对。改完的模型出了问题,你还能随时切回原始模型做对照实验。 - 权重加载时务必考虑设备问题。在 CPU 上加载权重再迁移到 GPU 可以减少显存占用,写法是
torch.load(..., map_location='cpu')。 - 训练时如果冻结了部分层,那些层的 BatchNorm 最好也固定为 eval 模式,否则统计量一直在更新,容易出现训练和推理表现不一致的问题。
7. 写在最后:一些个人经验
这次把这些内容整理出来的过程中,我回想自己最初接触 PyTorch 修改模型时的经历,踩得最深的坑其实就是“重命名层导致的 key 错位”。明明是同一个网络结构,只是把某个子模块单独拿出来定义成类了,结果权重的 key 挂载不上,让我白折腾了一整天。后来我把load_partial_state_dict的函数固定成自己的模板,每次都用它打印差集信息,才彻底避免了这类问题。
如果你也要做类似的事情,我这里有两个最实在的建议:
第一,改模型前一定花十分钟把原始模型的state_dict逐层打印一遍。很多问题不需要查文档就能解决——你只需要看清 key 的路径、了解每一层的形状,就能预估改完之后哪些 key 会失效。
第二,把修改后的模型固定成一个自定义类,而不是在加载模型后临时动手脚。这样做的好处是可以复用你的修改,也不容易在调试中丢失修改记录。代码的可读性提升了,将来换数据集、换任务时也能直接用同一套结构。
实际做迁移学习项目时,预训练模型的修改只是第一步,后续还需要考虑冻结策略、学习率、数据增强,这些都会影响最终效果。但结构修正确了,后面的调整才有意义。希望这篇内容能帮你绕过那些我踩过的坑,少走一些弯路。