这次聊一个 PyTorch 里看着基础、实操里天天遇见的操作:固定(freeze)住部分网络。做迁移学习、自监督训练、多阶段任务的同学基本都绕不过它,但一旦问深一点——冻结到底冻结了什么,PyTorch 底层是怎么处理的,哪些情况你的冻结动作实际上失效了——能一次说清楚的人并不多。这篇文章是我多次踩坑之后的完整整理,覆盖从原理到写代码再到排查冻结失效的整个链路,适用于 PyTorch 1.9 以上版本,普通训练、混合精度、多卡训练都能直接对上号。
1. 你真的需要冻结吗:三种典型场景与一个常见误区
1.1 迁移学习里冻结的真正目的
最常见的场景就是迁移学习。拿一个预先训练好的骨干网络,换掉最后的分类头,然后在新数据集上微调。很多人一上来就习惯性地把 backbone 全冻住,只训练新加的全连接层,觉得这样“又快又稳”。
其实在不大的新数据集上,这个做法确实能防止过拟合,因为低层卷积提取的是边缘、纹理、形状这类通用特征,这些特征在大多数视觉任务里是可复用的。如果全部参数一起微调,随机初始化的分类头在训练初期会产生很大的梯度,这些梯度回传到 backbone 会把预训练特征冲得七零八落。所以冻结浅层、只训深层和头部,本质上是把“预训练学到的通用表示”当作稳定的底座,在上面长出一小块任务专用的结构。
但如果你手里数据很充足,完全冻结反而会限制模型的表达能力。更合理的做法是“部分冻结 + 分层学习率”——底层冻结,高层用小学习率微调,头部用正常学习率训练。这个节奏我在本文最后一部分会专门讲。
1.2 对比学习、多阶段训练和受限资源场景
除了迁移学习,另一个高频场景是对比学习和自监督训练。典型做法是:先用无监督或大规模数据把 backbone 训练好,然后把 backbone 冻结,只训练一个很小的投影头(projection head)或线性分类头,用来做特征评估或下游任务适配。这种情况下冻结不是怕过拟合,而是因为 backbone 已经很强,再动它反而会让特征分布漂移。
还有一个很实际的原因:显存。冻结参数后,反向传播不会穿过被冻结的层,autograd 引擎就无需为这些层保存反向计算所需的中间激活值。前向计算仍然要做,但反向的存储和时间开销省下一大截。我在一个图像分类项目里试过,冻结了 ResNet 的前四层之后,单卡显存占用直接降了三分之一左右,训练速度也快了不少。算力紧张时,这种收益非常真实。
1.3 冻结不等于“前向不用算”
这里有一个新手很容易误解的地方:冻结(freeze)只影响反向传播,不影响前向计算。
requires_grad=False只是告诉 autograd“这个参数不需要梯度”,但 forward 时它依然参与所有计算。所以别指望冻结几层之后前向推理也会变快,前向计算量一点没省。省的是显存里保存的中间激活,以及反向传播的耗时。想省前向计算,只能靠减少网络层数、量化或者剪枝,那是另一个话题。
另一个误区是把“冻结参数”和“切换 eval 模式”混为一谈。冻结参数管的是参数是否更新,而model.train()/model.eval()管的是 BatchNorm 和 Dropout 的行为方式。两者是独立的。很多人只设置了requires_grad_(False)就以为万事大吉,结果 BatchNorm 的统计量照样在跑,这坑我后面专门讲。
2. requires_grad 与 autograd:冻结在底层做了什么
2.1 一个 Parameter 的四个字段
在 PyTorch 里,网络中的可训练参数都是nn.Parameter,它本质上是一个 Tensor,只是默认打开了requires_grad=True。每个 Tensor 上都有几个和自动求导相关的字段:data存放实际数值,grad存放梯度,requires_grad表示是否需要梯度,grad_fn记录它是通过什么运算产生的。
当你调用param.requires_grad_(False)时,你改的是requires_grad这个布尔量。别看只是一个布尔值,autograd 引擎在构建计算图时就是靠它决定要不要把当前节点纳入反向传播链路的。
2.2 计算图剪枝与叶节点
PyTorch 的反向传播是动态建图的。每次 forward 时,autograd 会把参与运算的 Tensor 和算子记录成一个图,然后 backward 从输出节点一路往回推梯度。
如果一个参数requires_grad=False,autograd 在 forward 阶段就会把它当成“常数”。反向传播推到这一层时,梯度不会继续往它前面的节点传。举个例子:ResNet 的 backbone 全冻结,头部的线性层需要梯度。backward 时梯度传入线性层的输入,也就是 backbone 的输出;因为 backbone 的所有参数都不需要梯度,autograd 在这里就会截断,backbone 内部的卷积、BatchNorm、激活函数等节点的梯度全部不需要计算。
这也就是为什么冻结能省显存:很多层的中间激活原本要留到 backward 时用,现在确认不需要梯度了,引擎就不会保留它们。
还有一个细节值得知道:requires_grad=False的张量通常可以看作“叶子节点”。对于叶子参数,你需要手动设置requires_grad;对于非叶子节点,即使你把它改成requires_grad=True,中间结果也不会凭空长出梯度来。所以实际使用中,我们通常是针对model.parameters()里的叶子参数操作。
2.3 为什么 requires_grad 管不到 BN 的统计量
这个太关键了,单独拆出来讲。
BatchNorm 层里有两类状态:一类是weight和bias,是nn.Parameter,会被requires_grad控制;另一类是running_mean和running_var,也就是滑动平均统计量,它们属于 buffer(缓冲区),不是 Parameter。buffer 默认不参与梯度计算,PyTorch 的model.parameters()也不会返回它们,所以你的冻结代码无论怎么写,都影响不到这些统计量的更新。
BatchNorm 在训练模式下的行为是:计算当前 batch 的均值和方差,用它来归一化输入,同时用滑动平均更新running_mean和running_var。这个更新过程发生在 forward 阶段,由 BatchNorm 模块内部完成,完全绕开 autograd。也就是说,即使你把这个 BN 模块里的weight和bias都冻结了,只要模型处于train()状态,它的统计量照样跟着当前 batch 的数据走。
这就是“冻结了但没完全冻结”的典型情况。后面我会给出完整的处理方案。
3. 冻结参数的四种写法,为什么要按场景选
3.1 最常用:遍历参数置 False,再定向解锁
直接看代码。
import torch.nn as nn def set_requires_grad(module: nn.Module, requires_grad: bool): for p in module.parameters(): p.requires_grad_(requires_grad) # 示例:冻结 backbone,保持新加的 head 可训练 model = create_model() set_requires_grad(model.backbone, False) set_requires_grad(model.head, True) # 如果 head 本来就是新加的,默认就是 True这是最直白、最不容易出错的写法。nn.Module本身没有requires_grad_()方法,只有nn.Parameter有,所以要遍历模块的parameters()逐个设置。封装成上面这个函数后,用起来就顺手多了。
有些人会写成for name, param in model.named_parameters(): param.requires_grad = False,效果一样。需要注意的是,named_parameters()默认不包含 buffer,所以 BN 的running_mean和running_var不会出现在这个循环里,这是预期行为。
3.2 名字匹配:按路径批量冻结
如果你只想冻结特定名字的层,可以用named_parameters()结合字符串匹配。
for name, param in model.named_parameters(): if name.startswith("backbone.layer1") or name.startswith("backbone.layer2"): param.requires_grad_(False) elif "classifier" in name: param.requires_grad_(True)named_parameters()打印出来的名字就是模块路径,比如backbone.layer1.0.conv1.weight,这样能精确到某个 block、某一次卷积,比整体冻结灵活得多。适合做消融实验,或者只冻结前几层、放开后面几层的情况。
这里我踩过一个坑:如果你用for name, param in model.named_parameters()判断完requires_grad后,又想用同样的name去getattr(model, name)取模块,会非常别扭。更干净的做法是设计一个前缀列表,按前缀判断,而不是按完整名字逐个判断。
3.3 优化器侧过滤:先冻结再建优化器
很多官方教程里推荐这种写法:
optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)用一个filter把requires_grad=True的参数挑出来,再传给优化器。这个写法非常常见,因为它把“哪些参数训练”的控制权完全交给了优化器构造这一步。
但它有个隐含前提:必须在构造优化器之前完成所有冻结动作。优化器在构造时会把参数引用放进param_groups,之后你再改requires_grad,优化器并不知道。
有人可能会问:优化器 step 的时候发现参数grad is None,不也会跳过吗?确实会跳过,但跳过的前提是grad为 None。如果一个参数在冻结前已经积累了梯度,之后又没做zero_grad(),这个旧梯度可能残留下来。冻结后它不再参与反向传播,grad不会被新梯度覆盖,优化器 step 时看到有旧梯度,照样更新。你明明冻结了参数,它的数值却还在变——这种诡异现象我遇到过不止一次。所以结论很简单:先冻结,再建优化器,别偷懒。
3.4 参数分组“软冻结”:lr=0 不是万能的
除了硬性把requires_grad关掉,还有人用参数分组的方式,把某些组的学习率设为 0,实现“软冻结”。
param_groups = [ {"params": model.backbone.parameters(), "lr": 0.0}, {"params": model.head.parameters(), "lr": 1e-3}, ] optimizer = AdamW(param_groups)这样 backbone 参数在优化器眼里依然是“可训练”的,只是每步的学习率是 0,所以数值不会变。但它和真正冻结有本质区别:requires_grad=True的参数会正常参与反向传播,正常产生梯度,autograd 不会剪枝,显存和计算上的节省完全不存在。
那这个写法有什么用?两个场景:一是你希望训练中途动态解冻某些层,用参数分组可以做到“只改 lr 不改 requires_grad”,免去重建优化器的麻烦;二是你希望某个模块参与梯度计算、但暂时不让它更新,比如某些带辅助损失的网络,先通过它反传梯度到前层,但参数本身先不动。
四种方式各有各的适用位置,我整理成了表格:
| 写法 | 控制维度 | 是否剪枝反向图 | 中途切换成本 | 典型场景 |
|---|---|---|---|---|
遍历置requires_grad=False | 参数级/模块级 | 是 | 需重建优化器 | 迁移学习、常规微调 |
| 按名字匹配冻结 | 参数路径级 | 是 | 需重建优化器 | 消融、部分层实验 |
优化器侧filter | 优化器构造期 | 间接生效 | 低 | 配合前两种写法收尾 |
参数分组设lr=0 | 优化器分组 | 否 | 低,改 lr 即可 | 动态解冻、分层训练 |
4. 冻结 BN 与 Dropout:最容易翻车的一层
4.1 统计量更新不受 requires_grad 控制
前面已经讲了原理:BN 的running_mean和running_var是 buffer,不参与model.parameters()。所以不管你用哪种方式冻结参数,只要把模型整体留在train()模式,冻结部分的 BN 统计量都会随着训练数据持续漂移。
举个实际例子。我用某个预训练模型做二分类微调,backbone 全部冻结,只训练最后两层。跑了两千步之后,发现 backbone 第一个 BN 层的running_mean和初始值已经明显不一样了。训练集分布跟预训练数据分布差异越大,漂移越严重。也就是说,预训练模型的统计信息被污染了,后面就算你想解冻或用它做推理,特征分布已经不是原来那个分布了。
4.2 冻结部分模块的正确姿势:按模块 eval()
处理办法是:把需要冻结的 BN 模块单独切到eval()状态,而不是对整个模型调用model.eval()。因为整个model.eval()会把头部、分类层这些还在训练的模块也一并切走,Dropout 也会关掉,有时候反而影响训练效果。
封装一个函数比较省心:
from torch.nn.modules.batchnorm import _BatchNorm def freeze_bn(module: nn.Module, freeze: bool = True): for m in module.modules(): if isinstance(m, _BatchNorm): m.eval() if freeze else m.train()用_BatchNorm这个基类判断,能同时覆盖BatchNorm1d、BatchNorm2d、BatchNorm3d和SyncBatchNorm,不用一个一个列。
同理,如果你的冻结部分里有 Dropout,也需要让这些模块进入eval()状态。原因很直观:Dropout 在训练模式下会随机丢弃特征,冻结层的前向输出会抖动,影响下游分类头接收到的特征稳定性。让冻结层保持确定性的输出,训练才会稳。
4.3 一个污染预训练统计量的复现与修复
我把这个坑整理成可复现的排查流程,方便大家对照:
- 断点检查:训练十步左右,拦住程序,打印 backbone 里任意一个 BN 模块的
running_mean变化量。 - 如果变化明显,说明该 BN 处于训练模式。你冻结的是权重,不是统计量。
- 修复:对冻结部分调用
freeze_bn(model.backbone),确保它们进入 eval 模式。 - 修复后再跑十步,
running_mean应该基本不变了。
这里再补充一个冷门但真实的细节:BatchNorm还有一个num_batches_tracked字段,只有启用track_running_stats=True时才会有。它记录当前统计量见过的 batch 数量,在running_var更新时用于无偏估计。有些分布式训练场景下,不同卡的num_batches_tracked同步不一致会导致 BN 统计行为有细微差别。如果项目对 BN 统计非常敏感,冻结时记得检查这个字段是否被意外更新。
5. 从零搭一个“冻结+微调”训练骨架,并验证冻结真的生效
5.1 完整骨架代码
把前面所有要点串起来,一个比较完整的骨架是这样的:
import torch import torch.nn as nn from torch.optim import AdamW def set_requires_grad(module: nn.Module, requires_grad: bool): for p in module.parameters(): p.requires_grad_(requires_grad) def freeze_bn(module: nn.Module, freeze: bool = True): from torch.nn.modules.batchnorm import _BatchNorm for m in module.modules(): if isinstance(m, _BatchNorm): m.eval() if freeze else m.train() # 1. 构建模型,更换任务头 model = torchvision.models.resnet50(pretrained=True) model.fc = nn.Linear(2048, num_classes) # 2. 冻结 backbone,保留 fc 可训练 set_requires_grad(model.backbone, False) set_requires_grad(model.fc, True) # 3. 冻结 backbone 内部所有 BN 的统计量更新 freeze_bn(model.backbone) # 4. 先冻结,后建优化器 optimizer = AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3, )注意第 4 步中的filter返回的是一个生成器,AdamW内部会把它转成列表,通常没问题。但如果后面要复用这个参数集合做梯度裁剪或者别的操作,建议顺手加个list(),避免生成器被迭代一次后变空。
5.2 三种验证手段:查 requires_grad、查 grad、看曲线
冻结完别急着跑,先验证一下你的冻结真的生效了。
第一种,直接检查requires_grad属性:
for name, p in model.named_parameters(): print(name, p.requires_grad)对照你的预期,确认冻结层全部是 False,训练层全部是 True。
第二种,训练一步之后打印梯度范数:
optimizer.zero_grad() loss = criterion(model(x), y) loss.backward() for name, p in model.named_parameters(): if p.requires_grad: grad_norm = p.grad.norm().item() if p.grad is not None else None print(f"{name}: grad_norm={grad_norm}") else: print(f"{name}: frozen, grad={None if p.grad is None else 'STALE!'}")如果某个requires_grad=False的参数打印出了STALE!,说明存在残留梯度,赶紧查一下是不是zero_grad没做好,或者冻结动作在优化器之后。
第三种,跑一个小规模实验看 loss 曲线。正确冻结时,loss 应该能稳步下降。如果整体 loss 完全不动,并且优化器里头参数一个都没有,PyTorch 会直接报 “optimizer got an empty parameter list”,这条报错其实就是“你冻过头了”的提示。
5.3 多卡和混合精度训练中的注意
切换到多卡训练时,冻结的时机更敏感。建议在DistributedDataParallel构造之前完成所有冻结决策。虽然 DDP 在反向传播时会自动跳过requires_grad=False参数的梯度同步,从而减少通信量,但如果你在训练中途改requires_grad,已构造好的 DDP 并不会平滑地适配这种变化,比较稳妥的做法是重新构造 DDP。
混合精度训练下要注意的是:autocast上下文中的前向计算仍然会让 BN 统计量正常更新,所以 AMP 并不能替你“冻结 BN”。如果用了torch.compile,动态改requires_grad还可能触发重新编译,建议在编译前就把冻结策略定好。
6. 冻结失效排查与渐进解冻的实战经验
6.1 场景一:冻结写在优化器之后
排查链路是这样的:检查代码里的执行顺序。如果你构造优化器的时候传了model.parameters(),之后又对其中某些参数调用了requires_grad_(False),优化器参数组里仍然留着这些参数。更麻烦的是残留梯度问题,冻结前的旧梯度会留在param.grad里,优化器 step 时照样更新。
我建议的排查顺序是:先看冻结代码是否在优化器之前,再看训练循环中zero_grad()是否在每次 loss 计算前执行。这两点都确认没问题,再往下查别的。
6.2 场景二:只关心参数,忽略了模块状态
这个场景就是前面讲的 BN 和 Dropout。参数确实不更新了,但模型内部的 buffer 统计量和随机行为还在变化。排查时除了打印requires_grad,还要打印额外的东西。
for name, m in model.named_modules(): if isinstance(m, _BatchNorm): print(name, m.training, m.running_mean[:3])如果m.training是 True,说明该 BN 还在训练模式,需要调用freeze_bn处理。
6.3 场景三:自定义模块里的“隐形参数”
有一种更隐蔽的情况:某些自定义模块在forward内部创建了临时权重,而不是用nn.Parameter。比如在forward里写w = torch.randn_like(x)然后直接传给卷积函数。这种临时张量既不在model.parameters()里,那你能不能用requires_grad冻结它。这类写法常见于一些简化的教学代码或快速原型,实际项目里如果出现训练行为诡异,先检查自定义层的forward里有没有这种“逃逸”的临时变量。
另外,用load_state_dict(strict=False)加载预训练权重后,如果网络里有些层改了名字,strict=False 只是跳过匹配不上的参数,并不会自动帮你决定哪些层需要冻结。排查时要确认目标层的参数真的在加载后的模型里,并且名字符合你的冻结判断条件。
6.4 渐进解冻与分层学习率:我自己的节奏
聊完排坑,最后分享我比较推荐的微调节奏:渐进解冻。
先只训练新加的头部,backbone 全部冻结。等头部的 loss 降到一个相对稳定的平台后,再解冻 backbone 最后的 stage(比如 ResNet 的 layer4),用骨干学习率的十分之一去微调。如果验证集指标没有恶化,再继续解冻更早的 stage。如果某一步解冻后指标不升反降,说明这一层及更前面的特征已经被预训练数据锻炼得很好了,新任务数据量不足以支持它们继续调整,那就停在这一步。
实现方式推荐参数分组,而不是反复改requires_grad:
param_groups = [ {"params": [p for p in model.backbone.parameters() if p.requires_grad], "lr": 1e-4}, {"params": model.fc.parameters(), "lr": 1e-3}, ] optimizer = AdamW(param_groups)解冻某个模块时,把它的requires_grad改回 True,然后重新构造优化器,或者直接把对应分组的lr从 0 调成目标值。我个人更喜欢前者,因为冻结/解冻本身是会改变反向传播链路的,重建优化器能让一切重新对齐,避免陈旧状态带来的隐患。
调试这段逻辑时,我强烈建议在训练脚本里加一个小工具函数,专门打印当前可训练参数的比例,以及每个分组的参数量和 lr。跑每个阶段之前先过一眼,能省下不少排查时间。
这些操作单独看都很琐碎,但组合起来就是一套可靠可复现的冻结微调流程。希望这份整理能让你少踩几次我踩过的坑。