1. 从一次梯度消失的排查说起
前阵子帮一位做模型微调的朋友看代码,他遇到一个很典型的问题:明明在训练脚本里把新加的分类头参数都设成了可训练,跑了几轮之后却发现这些参数的梯度始终是None,loss 也不怎么下降。我让他把模型结构打印出来一看,问题出在一个很隐蔽的地方——他在加载预训练权重之后,对整个模型做了一次model.requires_grad_(False),然后只对分类头重新打开了梯度,但分类头里有一层是复用的主干网络输出,那部分参数在冻结时被一起关掉了,后面又忘了单独打开。
这个坑其实非常常见。torch.Tensor.requires_grad_()这个带下划线的方法,看起来只是把requires_grad属性改一下,但它牵扯到自动微分图的构建、参数冻结与解冻、内存占用、甚至优化器的行为。很多人只记住了“冻结参数用requires_grad=False”,却忽略了它作用在Tensor上和在nn.Module上的区别,也忽略了原地操作带来的连锁反应。
这篇文章就围绕torch.Tensor.requires_grad_(requires_grad=True)这个具体方法展开,把它的语义、适用场景、和nn.Parameter的关系、在迁移学习和对抗训练里的实际用法、以及我踩过的那些坑,一次性讲清楚。不管你是刚接触 PyTorch 自动微分的新手,还是已经写过不少训练脚本的老手,只要你在做参数冻结、梯度裁剪、或者多阶段训练,这篇内容都能直接拿去对照排查。
2. requires_grad_ 到底改了什么
2.1 属性语义与原地操作的本质
先把这个方法的签名摆出来:
Tensor.requires_grad_(requires_grad=True) -> Tensor它接收一个布尔值,默认是True,作用是把当前张量的requires_grad属性设置为指定值,并且返回这个张量本身。注意关键词是“原地”(in-place),也就是说它不会创建新张量,而是直接修改调用它的那个张量的属性。
这里有个容易混淆的点:requires_grad是张量的一个属性,不是数据。它不改变张量里存的数值,只改变 PyTorch 在反向传播时是否要为这个张量计算梯度。你可以把它理解成给张量贴了一个标签:“我需要被追踪”。
为什么要有这个方法?因为requires_grad属性在张量创建之后,虽然可以直接赋值,比如x.requires_grad = True,但 PyTorch 官方更推荐用带下划线的方法形式。原因在于,直接赋值属性在某些版本和某些张量类型上可能不会触发内部的版本计数和依赖检查,而requires_grad_()是一个正式的原地操作接口,语义更明确,也更安全。
提示:对叶子张量(leaf tensor)直接设置
requires_grad是允许的,但对非叶子张量(即由运算产生的中间结果)设置requires_grad会报错,因为中间结果的梯度追踪状态由计算图决定,不能随意更改。这一点后面会详细说。
2.2 和 nn.Parameter 的关系
在nn.Module里,可训练参数通常不是裸的Tensor,而是nn.Parameter。nn.Parameter是Tensor的子类,它在创建时默认requires_grad=True。当你把一个nn.Parameter赋值给模块的属性时,它会自动注册到模块的parameters()列表里,优化器才能拿到它。
那么requires_grad_()能不能用在nn.Parameter上?可以,因为nn.Parameter继承自Tensor,方法都能用。但这里有个关键区别:对nn.Parameter调用requires_grad_(False),只是把这个参数的梯度追踪关掉,它仍然留在parameters()列表里。优化器如果已经用model.parameters()初始化了,它依然会拿到这个参数,只是这个参数的.grad会一直是None,优化器更新时相当于跳过它(不同优化器行为略有差异,后面会讲)。
这就引出一个实操中很常见的误区:很多人以为requires_grad_(False)就等于把参数从模型里“移除”了,其实没有。参数还在,只是不参与梯度计算。如果你想让优化器完全不碰它,要么在初始化优化器之前就冻结好,要么用参数组(param groups)显式过滤。
2.3 叶子张量与非叶子张量的分界线
这是理解requires_grad_()的核心。PyTorch 把张量分成两类:
- 叶子张量(leaf tensor):由用户直接创建的张量,比如
torch.randn(3, 4),或者nn.Parameter。叶子张量的requires_grad可以随时改。 - 非叶子张量(non-leaf tensor):由运算产生的张量,比如
y = x * 2里的y。非叶子张量的requires_grad由输入张量和运算类型决定,不能手动改。
如果你对一个非叶子张量调用requires_grad_(),会直接抛异常:
x = torch.randn(3, 4, requires_grad=True) y = x * 2 y.requires_grad_(False) # RuntimeError: you can only change requires_grad flags of leaf variables.这个报错信息很直白:只能改叶子变量的requires_grad标志。为什么这么设计?因为非叶子张量的梯度追踪状态是计算图的一部分,如果允许随意改,反向传播的链条就断了,自动微分会变得不可预测。
那实际中怎么处理?如果你确实想让某个中间结果不参与梯度,应该用detach(),它会返回一个和原张量共享数据但脱离计算图的新张量,新张量是叶子,可以自由设置requires_grad。
y = (x * 2).detach() y.requires_grad_(True) # 合法,因为 y 现在是叶子注意:
detach()和requires_grad_(False)不是一回事。detach()返回新张量,原张量还在图里;requires_grad_(False)是原地改属性,只对叶子有效。混用这两个概念是新手最容易犯的错之一。
3. 参数冻结与解冻的完整实操
3.1 迁移学习中的标准冻结流程
迁移学习是最常用到requires_grad_()的场景。假设你加载了一个预训练的主干网络,想冻结它的全部参数,只训练新加的分类头。标准做法是这样的:
import torch import torch.nn as nn class MyModel(nn.Module): def __init__(self, backbone, num_classes): super().__init__() self.backbone = backbone self.classifier = nn.Linear(backbone.out_features, num_classes) def forward(self, x): features = self.backbone(x) return self.classifier(features) model = MyModel(pretrained_backbone, num_classes=10) # 第一步:冻结主干 for param in model.backbone.parameters(): param.requires_grad_(False) # 第二步:确认分类头是可训练的 for param in model.classifier.parameters(): param.requires_grad_(True) # 第三步:只把需要梯度的参数交给优化器 optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3 )这里有几个细节值得展开。第一,冻结操作要在初始化优化器之前做,或者像上面这样用filter过滤。如果你先初始化了优化器,再把参数冻结,优化器内部可能已经保存了这些参数的引用和状态(比如 Adam 的动量),虽然后续更新时因为.grad是None而跳过,但状态里可能残留旧值,切换阶段时容易出问题。
第二,filter(lambda p: p.requires_grad, model.parameters())这个写法比直接传model.parameters()更稳妥。它确保优化器只管理真正需要更新的参数,减少不必要的内存开销,也避免了一些优化器在遇到grad=None时的边界行为差异。
第三,冻结之后建议打印一下确认:
for name, param in model.named_parameters(): print(name, param.requires_grad)这个习惯能帮你快速发现“以为冻结了其实没冻”或者“以为解冻了其实还冻着”的问题。我见过太多因为模块嵌套层级深、名字对不上而导致的冻结失效。
3.2 分阶段解冻的策略与代码
微调任务里,常见做法不是一上来就全解冻,而是分阶段解冻。比如先只训练分类头几轮,让新加的层先适应,然后再逐步解冻主干的后几层。这样做的好处是避免随机初始化的分类头在早期产生大梯度,把预训练好的主干权重带偏。
分阶段解冻的代码可以这样写:
def unfreeze_layers(model, num_layers): """解冻主干网络最后 num_layers 层""" backbone_layers = list(model.backbone.children()) for layer in backbone_layers[-num_layers:]: for param in layer.parameters(): param.requires_grad_(True) # 第一阶段:只训练分类头 for param in model.backbone.parameters(): param.requires_grad_(False) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3 ) train(model, optimizer, epochs=3) # 第二阶段:解冻最后两层,降低学习率 unfreeze_layers(model, num_layers=2) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4 ) train(model, optimizer, epochs=5)注意第二阶段我重新创建了优化器,并且把学习率调低了。为什么?因为解冻的层是预训练权重,已经比较好了,用大学习率容易破坏。重新创建优化器是为了让新加入的参数有干净的动量状态,避免继承第一阶段不存在的历史。
提示:如果你不想重建优化器,也可以用参数组的方式动态调整。但重建更简单直接,在阶段切换不频繁的场景下完全够用。
3.3 冻结 embedding 层的特殊处理
Embedding 层在 NLP 模型里很常见,冻结它的时候有个小细节。nn.Embedding的权重是一个nn.Parameter,形状是(num_embeddings, embedding_dim)。对它调用requires_grad_(False)是有效的,但如果你用了sparse=True的 embedding,梯度是稀疏的,某些优化器对稀疏梯度的支持和稠密梯度不同。
另外,如果你冻结了 embedding 但没冻结后面的层,反向传播时梯度传到 embedding 输出就停了,不会继续往 embedding 权重传。这是符合预期的。但要注意,如果 embedding 的输出被detach()过,那后面的层也拿不到梯度,整个链条就断了。排查时可以用torch.autograd.grad手动检查某个中间张量的梯度是否存在。
4. 对抗训练与梯度操作的配合
4.1 对抗样本生成中的 requires_grad 切换
对抗训练(比如 FGSM、PGD)需要在输入上计算梯度,而不是在参数上。这时候requires_grad_()就派上用场了。标准流程是:先把模型参数冻结,然后对输入张量打开梯度,前向传播后取输入的梯度,用这个梯度构造对抗样本。
def fgsm_attack(model, x, y, epsilon): # 冻结模型参数,避免污染 for param in model.parameters(): param.requires_grad_(False) # 对输入打开梯度 x_adv = x.clone().detach() x_adv.requires_grad_(True) output = model(x_adv) loss = nn.CrossEntropyLoss()(output, y) loss.backward() # 用输入的梯度构造对抗样本 with torch.no_grad(): x_adv = x_adv + epsilon * x_adv.grad.sign() x_adv = torch.clamp(x_adv, 0, 1) # 恢复模型参数的梯度状态 for param in model.parameters(): param.requires_grad_(True) return x_adv.detach()这段代码里有几个关键点。第一,x.clone().detach()先复制一份输入并脱离原计算图,然后requires_grad_(True)让它成为新的叶子张量。如果直接对原始x调requires_grad_(True),而x本身是数据加载器出来的、没有梯度的张量,那也可以,但为了不污染原始数据,克隆更安全。
第二,攻击结束后要把模型参数的requires_grad恢复。如果不恢复,后续正常训练时参数就不更新了,这是一个非常隐蔽的 bug。我建议把恢复操作放在finally块里,或者用上下文管理器封装,避免异常路径下忘记恢复。
第三,x_adv.grad在loss.backward()之后才有值,而且它和x_adv共享内存。构造对抗样本时要用torch.no_grad()包住,否则这个加法操作会被记录进计算图,导致后续反向传播出错。
4.2 梯度惩罚中的 detach 与 requires_grad
有些训练目标里包含梯度惩罚项,比如 WGAN-GP 里的梯度范数约束。这时候需要对输入求梯度,但又不希望这个梯度参与二阶求导。做法是先用requires_grad_(True)打开输入梯度,算出惩罚项后,对惩罚项里的张量做detach(),再参与总 loss。
def gradient_penalty(model, real, fake): alpha = torch.rand(real.size(0), 1, 1, 1) interpolated = (alpha * real + (1 - alpha) * fake).requires_grad_(True) output = model(interpolated) grad = torch.autograd.grad( outputs=output, inputs=interpolated, grad_outputs=torch.ones_like(output), create_graph=True, retain_graph=True )[0] grad_norm = grad.view(grad.size(0), -1).norm(2, dim=1) penalty = ((grad_norm - 1) ** 2).mean() return penalty这里create_graph=True是为了让惩罚项本身可导,从而参与参数更新。而interpolated是通过运算得到的,不是叶子,所以不能直接对它调requires_grad_()——等等,这里其实可以,因为(alpha * real + (1 - alpha) * fake)的结果如果real和fake都不需要梯度,那这个结果就是叶子。但如果real或fake需要梯度,结果就是非叶子,调requires_grad_()会报错。稳妥做法是先detach()再打开:
interpolated = (alpha * real + (1 - alpha) * fake).detach().requires_grad_(True)这个写法在 WGAN-GP 的实现里非常常见,也是我推荐的标准写法。
5. 常见报错与排查速查
5.1 典型错误信息对照表
| 报错信息 | 触发原因 | 解决方法 |
|---|---|---|
RuntimeError: you can only change requires_grad flags of leaf variables | 对非叶子张量调用requires_grad_() | 先用detach()得到叶子张量,再设置 |
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn | 反向传播时计算图里没有需要梯度的张量 | 检查是否所有参数都被冻结,或输入未打开梯度 |
UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed | 访问了非叶子张量的.grad | 用retain_grad()或改为访问叶子张量 |
| 优化器更新后参数不变 | 参数requires_grad=False或未加入优化器 | 打印param.requires_grad和优化器参数组确认 |
| 冻结后显存没降 | 冻结只影响梯度计算,不影响前向激活 | 冻结需配合torch.no_grad()才能省显存 |
5.2 冻结了但显存没降的排查
这是一个高频疑问。很多人冻结了主干参数,发现显存占用几乎没变。原因在于,requires_grad_(False)只是不计算这些参数的梯度,但前向传播的中间激活值仍然存在,因为反向传播到冻结层就停了,但冻结层之前的激活(如果还有可训练层)仍然需要保留。
真正省显存的做法是在冻结层的前向传播外面包torch.no_grad():
def forward(self, x): with torch.no_grad(): features = self.backbone(x) return self.classifier(features)这样主干的前向不构建计算图,中间激活不保存,显存能明显下降。但注意,这样做之后,即使你后来把主干参数解冻,这部分前向也不会产生梯度,需要把no_grad去掉。所以冻结策略要和前向代码保持一致。
5.3 多卡训练下的同步问题
如果你用DataParallel或DistributedDataParallel,冻结操作要在包装模型之前做,或者确保每个进程都执行了冻结。DataParallel会在每个前向时复制模型到各卡,如果冻结状态没同步,可能出现某些卡上参数可训练、某些卡上不可训练的情况,导致梯度不一致。
稳妥做法是在模型定义和冻结完成后,再包DistributedDataParallel。另外,requires_grad_(False)的参数在 DDP 里不会参与梯度同步,这能减少通信量,算是一个额外收益。
6. 几个我踩过的坑和实操建议
第一个坑是关于nn.Parameter的默认行为。nn.Parameter创建时默认requires_grad=True,但如果你用nn.Parameter(data, requires_grad=False)创建,它仍然是参数,只是不追踪梯度。有些人以为这样创建就不会被优化器拿到,其实还是会。要彻底排除,得在优化器初始化时过滤。
第二个坑是requires_grad_()的返回值。它返回的是原张量本身,所以x = x.requires_grad_(True)和x.requires_grad_(True)效果一样,但前者容易让人误以为创建了新张量。我建议直接调用不赋值,语义更清晰。
第三个坑是关于torch.no_grad()和requires_grad_(False)的混用。在推理阶段,两者都能省显存,但torch.no_grad()是上下文管理器,作用范围明确;requires_grad_(False)是持久状态,容易忘记恢复。推理时优先用torch.no_grad(),训练中需要长期冻结某部分参数时才用requires_grad_(False)。
第四个坑是保存和加载模型时的状态。requires_grad状态不会被state_dict()保存,所以加载模型后需要重新设置冻结。如果你依赖冻结状态做推理,记得在加载后重新执行冻结逻辑,或者把冻结逻辑写进模型的__init__里。
最后一个建议:在训练脚本里加一个检查函数,每个 epoch 开始时打印一次各模块的requires_grad状态。这个开销很小,但能帮你及早发现冻结失效或误冻结的问题。我现在的项目里都会保留这个检查,尤其是在做多阶段微调的时候,它救过我好几次。
提示:如果你用的是
torch.compile或某些图优化工具,requires_grad_()的原地修改可能会触发重新编译。在编译前确定好冻结策略,避免训练中途频繁切换,否则编译缓存会失效,性能反而下降。