PyTorch参数冻结完全指南：迁移学习微调中如何正确固定模型层

最近一个图像分类的迁移学习项目把我折腾得够呛。预训练模型在源任务上表现不错但拿到新数据集之后我顺手把整个网络扔进训练循环结果发现backbone的参数被疯狂微调训练速度慢得离谱而且在新数据集只有几千张图片的情况下很快就把预训练学到的通用特征搞坏了。直到我把冻结、固定部分参数这件事彻底搞懂问题才真正解决。这篇内容就围绕PyTorch里如何冻结、固定部分参数来展开把原理、写法、坑和进阶玩法一次性说清楚。如果你正在做迁移学习微调、多任务训练或者想把某个特征提取器固定下来只训练新的分类头这篇文章应该能让你少踩几个坑。1. 为什么冻结和固定不完全是一回事参数更新的判定链先说结论在PyTorch里一个参数最终会不会被更新取决于一条完整的判定链而不是单单看某个函数或某个参数名。1.1 训练循环里的三段式一个标准的PyTorch训练循环可以拆成这样for x, y in dataloader: optimizer.zero_grad() # 清空上一次的梯度 loss criterion(model(x), y) loss.backward() # 反向传播计算每个requires_gradTrue参数的梯度 optimizer.step() # 根据梯度更新参数参数会不会更新关键看optimizer.step()这一行。而step()遍历的是优化器在构造时保存下来的参数列表。这个列表里的每个参数只要grad不是None就会被按优化器规则更新。所以逻辑链其实很清晰反向传播时只有requires_gradTrue的叶子参数才会被计算梯度。优化器构造时只接收它看到的参数。step()时遍历优化器的参数列表只有grad不为None的参数才会被更新。换句话说想冻结一个参数有两条路要么让它没有梯度设requires_gradFalse要么让它不出现在优化器里。两条路都行但效果和副作用不一样后面细说。1.2 不更新和固定住差在哪不更新指的是参数值在训练中保持不变这通常就是我们在微调时想要的。但固定住还有另一层含义某些和模型行为强关联的内部状态比如BatchNorm层的running_mean和running_var它们不是Parameter而是buffer。即使你把所有requires_grad都设成False只要模型处于train()模式前向传播时BN仍然会在原地更新这些统计量。也就是说参数没动但模型内部状态还在变。这是很多人踩坑的地方冻结了参数却没冻结行为。所以我在实际项目中一般把冻结拆成三个独立开关来理解。开关控制对象手段参数是否计算梯度weight、bias等Parameterrequires_gradFalse参数是否出现在优化器优化器持有的参数列表filter或参数组模块的行为模式BN统计量更新、Dropout开关model.eval()/module.train()这三者要配合使用才叫完整的冻结。只做其中一两项很容易出现参数没变但loss乱跳或者backbone参数纹丝不动BN统计量却已经面目全非的尴尬局面。2. 标准冻结三件套requires_grad、优化器过滤、模型模式2.1 第一件设requires_gradFalse最常见的场景是加载预训练模型替换分类头然后冻结整个backbone。import torch import torchvision.models as models model models.resnet18(pretrainedTrue) # 冻结所有层 for name, param in model.named_parameters(): param.requires_grad False # 替换分类头 model.fc torch.nn.Linear(model.fc.in_features, 10)注意named_parameters()会递归遍历所有子模块包括你刚替换的fc层。所以如果先冻结全部再替换fc新层默认requires_gradTrue正好只训练分类头。如果你用的不是torchvision的模型是自己定义的专业网络同理。关键是循环所有parameters()或是named_parameters()把需要固定的层设为False。2.2 第二件过滤优化器参数列表在很多教程里冻结参数后直接用model.parameters()构造优化器也能跑通。因为requires_gradFalse的参数在backward后grad是Nonestep()会跳过它们。这是对的但不是最稳妥的。更稳妥的写法是显式过滤optimizer torch.optim.SGD( filter(lambda p: p.requires_grad, model.parameters()), lr0.01 )为什么我要强调这一步有两个原因。第一个原因是性能。如果优化器里挂着几百个参数的引用虽然不更新但每次step()仍然要遍历一遍判断grad is None。在小模型上无所谓但大模型训练时这个遍历开销和潜在的状态管理成本是实实在在的。第二个原因是后续操作的确定性。如果你把requires_grad设为False后不小心在某个地方又把它设回了True那么优化器里的这个参数会立刻开始更新——哪怕你本意是永远冻结。显式过滤后哪怕你忘了某处把它改回True优化器里根本没有它也不会误更新。等于加了一层保险。2.3 第三件管理模型模式这一步尤其在含BN和Dropout的模型里必须做。大多数教程只讲参数冻结不讲eval和train的模式管理导致不少人看到参数没更新但验证集效果越来越差的诡异现象。标准做法是model.eval() # 全局先切到eval model.fc.train() # 只把要训练的那部分切回train先eval()再局部train()的顺序至关重要。如果反过来先对某个子模块train()再调用model.eval()后者会递归把所有子模块都切回eval你前面的设置就白做了。为什么需要这样因为冻结backbone时我们通常希望它完完全全像一个固定的特征提取器BN统计量不更新Dropout直接关闭。而新的分类头则要保持正常的训练行为比如Dropout要继续生效。这时候必须让两个部分处于不同的模式。我把这三件套封装成一个简单的小函数每次换任务都能复用def freeze_module(module: torch.nn.Module, freeze: bool True, keep_train: bool False): for param in module.parameters(): param.requires_grad not freeze if freeze and not keep_train: module.eval() elif not freeze: module.train() return modulekeep_train参数解决的是前文说的情况有时候我只想冻结参数更新但希望BN统计量继续适应新数据分布。这时候freezeTrue但keep_trainTrue参数不更新BN还在train模式下更新统计量。注意这只是参数冻结不是行为冻结两种需求要区分开。3. 最容易翻车的层BN的统计量更新和Dropout的开关3.1 BN的running_mean/running_var到底谁在更新BatchNorm2d这种模块里有两类东西weight和bias是Parameter受优化器控制。running_mean和running_var是buffer存储在state_dict里但不被优化器更新。那running_mean和running_var什么时候变答案是在前向传播时由F.batch_norm根据当前batch统计量原地更新条件是模块处于train()模式。很多人以为设了requires_gradFalseBN就完全不动了这是误解。你设的只是weight和bias的梯度running_mean照样每步都在滑动平均。所以在冻结backbone的场景里如果你希望整个backbone完全固定必须让backbone处于eval()模式。比如model.backbone.eval() model.fc.train()这样backbone的BN就以推理模式运行用历史统计量做归一化不再原地更新。3.2 两个反向的典型案例我在实际项目中遇到过两种截然相反的诉求一直留着当教训。第一种诉求是我想让backbone参数冻结但BN统计量最好还是适应新数据。这种场景其实挺常见预训练模型来自A数据集目标数据集B的分布略有偏移。这时候你可以让backbone的weight/bias不更新但保持train()模式让running_mean/running_var继续滑动。注意在PyTorch的F.batch_norm实现里只要trainingTrue就更新running stats和参数requires_grad无关。第二种诉求是整个backbone完全变成死的特征提取器连BN统计量也不要变。那就必须eval()。但这里有个新坑如果你之后要对全模型model.train()恢复训练backbone也会被切回train模式BN又开始更新了。所以涉及冻结一部分、训练一部分的代码里model.train()这种全局调用要非常谨慎最好手动指定各子模块模式。3.3 Dropout的开关逻辑Dropout的原理是训练时随机丢弃一部分神经元的输出推理时全部保留并按概率缩放。但它的实现依赖模块模式train()模式下生效eval()模式下直接透传。如果你用model.eval()冻结了整个backbone而新的分类头里有Dropout层那么分类头的Dropout也一起被关掉了。这对某些任务可能影响收敛尤其在少样本情况下训练时没有随机性模型更容易过拟合。正确的处理方式和我前面写的一样先model.eval()再手动把需要训练的部分.train()。千万别反过来。3.4 Batch Size等于1的隐藏陷阱还有一个和BN强相关的坑很多人在冻结后把batch_size调成1来调试结果直接报错ValueError: Expected more than 1 value per channel when training, got input size ...原因BN在train模式且batch size为1时无法计算批量统计量的方差。哪怕backbone参数被冻结只要backbone的BN还处于train()模式这个错就会炸出来。我一开始也以为是模型哪里写错了查了半天最后发现就是model.eval()没设。所以在冻结场景里如果要batch_size1调试要么把backbone的BN切到eval要么干脆整个模型先eval。这类边界条件在长训任务中很常见不值得半夜在线排查。4. 细粒度控制部分层可学、中途解冻与分组学习率4.1 只让最后几层可学习有些任务里整个backbone的浅层特征是有用的深层特征和目标任务不太匹配。这时候可以只解冻最后几个block或最后一层。用named_parameters()配合名字判断即可# 只让layer4和fc可以学习 for name, param in model.named_parameters(): if name.startswith(layer4) or name.startswith(fc): param.requires_grad True else: param.requires_grad False如果不想依赖硬编码名字也可以用子模块路径来控制for param in model.layer1.parameters(): param.requires_grad False for param in model.layer4.parameters(): param.requires_grad True注意named_parameters()返回时每一层的参数都是模块里的叶子参数递归设置requires_grad是可以直接生效的。之后再backward()只有requires_gradTrue的参数路径会被记录。4.2 两阶段训练先冻结训练head再解冻整体微调这是迁移学习的经典套路。第一阶段backbone冻结只训练新分类头第二阶段把backbone解冻用很小的学习率整体微调。# 阶段一 for param in model.backbone.parameters(): param.requires_grad False model.backbone.eval() model.fc.train() optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3 ) # 训练若干个epoch... # 阶段二解冻 for param in model.backbone.parameters(): param.requires_grad True model.backbone.train() optimizer torch.optim.AdamW( model.parameters(), lr5e-5 ) # 继续微调...这里我强烈建议在阶段一切换到阶段二时重新构建优化器。因为阶段一的优化器状态都在fc的参数上解冻backbone之后同一个优化器内部旧有的state并不会自动适配到新参数上。而且如果你用的是Adam这类带动量项的优化器长期使用同一个优化器容易让解冻后的更新步长变得很怪不如直接重建一个干净的优化器。4.3 分组学习率backbone和分类头用不同学习率很多老练的人不做显式二阶段而是直接用参数组让backbone学习率小一点、head学习率大一点一次训练搞定。optimizer torch.optim.AdamW([ {params: model.backbone.parameters(), lr: 1e-4}, {params: model.fc.parameters(), lr: 1e-3}, ])这样在一个训练循环里backbone更新得保守head更新得激进。等价于软冻结——不是完全不更新而是更新幅度很小。这个技巧在数据量中等、预训练特征非常合理时特别好用。它的优势在于backbone不完全是死的它还能根据目标任务的信号做微调同时因为学习率低不容易破坏预训练特征。4.4 动态调整学习率的坑如果你使用了torch.optim.lr_scheduler.StepLR这类调度器它会遍历优化器的所有param_groups统一按比例衰减。这个行为在单参数组时没问题但分组学习率下就不一定符合你的预期。比如backbone初始lr1e-4head初始lr1e-3调度器每N个epoch把lr乘0.1那两组学习率会同时缩小10倍两者相对比例不变。这通常没问题。但如果你想要backbone完全不衰减、head衰减就需要自己写自定义调度器或在每个epoch手动修改param_groups里的lr。我建议用MultiStepLR这类现成的就行省心。真遇到特殊需求就每个epoch循环里手动写# 手工调整某一组的learning rate optimizer.param_groups[0][lr] 5e-5不要嫌手动简单直接不容易出错。5. 冻结后的自检清单与三个容易忽略的小坑5.1 训练前先跑一遍参数状态报告冻结操作写完后不能直接闷头训练。我会先跑一个自检函数打印出可学习的参数数量再检查目标层是否真的被冻结了。这个习惯帮我救回了好几次配错了冻结范围的尴尬。def parameter_report(model): total 0 trainable 0 frozen_layers [] trainable_layers [] for name, param in model.named_parameters(): n param.numel() total n if param.requires_grad: trainable n trainable_layers.append(name) else: frozen_layers.append(name) print(fTotal params: {total:,}) print(fTrainable params: {trainable:,} ({trainable / total:.2%})) print(Frozen layers:, frozen_layers[:5], ...) print(Trainable layers:, trainable_layers[:5], ...)每次换实验配置跑一遍这个报告几秒钟的事能让你一眼看清有没有把关键的层错误冻结。5.2 检查优化器里到底有哪些参数模型参数报告只能说明requires_grad的情况优化器里实际持有的参数不一定和它一致。尤其当你用filter(lambda p: p.requires_grad, ...)构造优化器之后如果又动态解冻了某些层优化器依然只持有构造时的参数列表并不会自动更新。所以我会额外写一个小工具来核对优化器的参数数量和模型的可学习参数数量是否一致def optimizer_report(optimizer): count sum( p.numel() for group in optimizer.param_groups for p in group[params] ) print(fOptimizer holds {count:,} params)不一致时大概率是你构建优化器后修改了requires_grad但没有重建优化器。5.3 三个容易忽略的小坑第一个坑model.train()把一切打回原形。我在多阶段训练时经常写完阶段二就调用一次model.train()结果backbone的BN统计量再次开始更新。后来我养成了习惯涉及冻结的代码一律不用全局train()只对需要训练的子模块逐个调用.train()或者写一个set_mode函数统一管理。第二个坑save和load时buffer的状态。state_dict里会包含BN的running_mean/running_var这些buffer。如果你在冻结前保存了模型解冻后再loadbuffer会按state_dict恢复。这本身没问题但如果你在冻结期间期望backbone的BN统计量保持初始值而模型又跑过几个epoch这时候保存的统计量其实已经变了下次load回来就是已经被动过的统计量。解决方式很简单冻结backbone时干脆把它的eval()模式保持住不让统计量变化。第三个坑冻结后hook拿到的梯度是None。很多人用register_forward_hook或register_full_backward_hook做特征可视化或梯度分析结果冻结参数后在hook里直接访问grad会得到None而非零向量导致奇怪的报错或者误判。正确做法是先判断if param.grad is not None: # 处理梯度对param.grad做判空处理别假定它存在。5.4 一个关于混合精度训练的小提醒如果用torch.cuda.amp做混合精度训练冻结层不会有特殊问题但要注意如果backbone完全冻结且处于eval()模式它的参数在前向中不需要梯度。此时PyTorch的autograd会自动优化不再为这些参数构建梯度图计算开销会小很多。如果你让backbone保留requires_gradFalse但又处于train()模式那么BN的running stats更新依然会发生并且反向传播的计算图会包含backbone直到BN统计量的路径存在一定的额外开销与潜在数值风险。所以从性能角度我遇到这类任务时都会尽量让冻结的部分保持eval()这既是行为固定也能让计算图更简洁。最后说说我现在的习惯。做任何迁移学习项目第一件事永远是写一份参数状态报告加一个模块模式管理函数把冻结哪个部分、哪些层要训、BN和Dropout的模式怎么设一次性管清楚。之后再谈调学习率、解冻微调这些进阶操作。这个流程看起来多花了几分钟但省掉了后面大量排查诡异训练曲线的时间。希望你也能在自己的任务里把这些坑提前绕开把时间花在真正有价值的调参和实验设计上。