PyTorch requires_grad_() 深度解析：梯度冻结与解冻的陷阱与最佳实践

1. 从一次“诡异”的梯度消失说起最近在复现一个经典的图像分类网络时我遇到了一个让人挠头的现象模型在训练初期损失值纹丝不动准确率也卡在随机猜测的水平。我检查了数据加载、损失函数、优化器甚至怀疑是不是学习率设得太小。一通排查下来问题竟然出在一个看似不起眼的地方——我在加载预训练模型权重后为了“保护”某些层不被更新手动将它们的requires_grad属性设为了False。后来当我需要重新微调这些层时我理所当然地使用了tensor.requires_grad_(True)来“激活”梯度计算。然而训练依然没有起色。这个经历让我意识到很多朋友对 PyTorch 中requires_grad和requires_grad_()的理解可能还停留在“开/关梯度”的层面。实际上它们的行为远比一个简单的布尔开关要微妙尤其是在涉及计算图构建、参数冻结与解冻、以及模型部署等场景时理解其深层机制至关重要。今天我们就来彻底拆解requires_grad_(True)这个操作看看它背后到底发生了什么以及我踩过的那个坑究竟是怎么形成的。简单来说requires_grad是 PyTorchTensor的一个属性它决定了在反向传播过程中是否需要计算该张量相对于某个标量损失函数的梯度。而requires_grad_()是一个就地in-place方法用于修改这个属性。但“设置为 True”并不总是意味着“梯度开始流动”其生效条件和影响范围是理解 PyTorch 自动微分Autograd系统的关键一环。2. requires_grad 与 Autograd 计算图不仅仅是属性开关要理解requires_grad_(True)我们必须先回到 PyTorch 自动微分的核心——动态计算图。当你对requires_gradTrue的张量进行操作时PyTorch 会跟踪所有操作构建一个由Function节点组成的有向无环图DAG。这个图记录了从输入到输出的完整计算路径。2.1 计算图的构建与梯度流假设我们有两个张量import torch a torch.tensor([1.0, 2.0], requires_gradTrue) b torch.tensor([3.0, 4.0], requires_gradTrue)此时a和b是计算图中的“叶子节点”leaf tensor它们是由用户直接创建的并且显式要求计算梯度。当我们执行运算c a * b # 逐元素相乘 d c.sum() # 求和得到一个标量PyTorch 会动态地构建计算图。c是MulBackward函数的输出d是SumBackward函数的输出。a和b是这些函数的输入。当我们调用d.backward()时Autograd 引擎会沿着这个图反向执行调用每个Function节点中定义的backward()方法将梯度从d一路传播回a和b。这里的关键在于计算图是在前向传播执行运算时即时构建的。一个张量的requires_grad属性决定了它是否会被记录到计算图中以及后续的运算是否会以它为基础继续构建需要梯度的子图。2.2 requires_grad_(False) 的真正含义从图中“剥离”当我们对一个叶子张量比如模型参数param调用param.requires_grad_(False)时我们做了两件事将张量本身的requires_grad属性设置为False。更重要的是这向 Autograd 引擎发出信号在未来的前向传播中任何以该张量为起点的运算其产生的中间张量默认都不会被纳入需要梯度的计算图。例如param torch.randn(3, 3, requires_gradTrue) # 初始需要梯度 param.requires_grad_(False) # 冻结该参数 # 前向传播 output model(input) # 假设model使用了param loss criterion(output, target) loss.backward() # 此时param.grad 将为 None因为计算图在param处“断开”了在这个例子中param的梯度不会被计算也不会被优化器更新。这是模型微调中冻结底层特征的常用技巧。2.3 requires_grad_(True) 的生效条件时机至关重要现在来到核心问题如果之后我想解冻这个参数调用param.requires_grad_(True)会发生什么答案是它只对未来的计算有效。requires_grad_(True)仅仅修改了张量自身的属性。它不会追溯性地修改已经构建好的计算图。之前因为param.requires_grad为False而导致没有记录梯度计算路径的运算其历史不会被改变。这就是我踩坑的原因。我的操作顺序是加载预训练模型所有参数requires_gradTrue。冻结部分层layer.parameters()遍历并设置param.requires_grad_(False)。训练几个 epoch冻结层无梯度。想要解冻这些层于是调用param.requires_grad_(True)。继续训练发现解冻的层仍然没有梯度更新。问题根源在第4步之后我直接开始了新的训练循环。但是优化器如torch.optim.SGD或Adam在初始化时会捕获所有requires_gradTrue的参数。当我第2步冻结参数后这些参数从优化器的param_groups中被“移除”严格说是优化器不再更新它们但PyTorch的优化器设计是仍然持有引用只是忽略grad为None的参数。当我第4步将requires_grad改回True后优化器并不知道这一变化它内部维护的参数列表和状态如动量并没有自动更新。所以正确的解冻流程应该是# 1. 设置参数需要梯度 for param in layer.parameters(): param.requires_grad_(True) # 2. 关键步骤将解冻的参数重新添加到优化器中或者创建一个新的优化器。 # 更简单通用的做法是在修改requires_grad后重新定义优化器只传入requires_gradTrue的参数。 optimizer torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.001)如果不重新创建或更新优化器即使梯度计算出来了param.grad不为None优化器的step()函数也可能不会正确地更新这些参数因为它的内部状态没有为这些“新加入”的可训练参数初始化相应的缓存如Adam的动量缓存。3. 深入辨析requires_grad_()、detach() 与 no_grad() 的边界在控制梯度流时我们还有detach()和torch.no_grad()这两个常用工具。它们与requires_grad_()有何区别理解这个能帮你更精准地控制模型行为。3.1 detach()创建无梯度的“副本”切断历史tensor.detach()会返回一个新的张量这个新张量与原始张量共享底层数据存储但它是一个新的叶子节点并且requires_gradFalse。最关键的是它从当前计算图中被“分离”出来其之前的运算历史被完全切断。x torch.tensor([1.0], requires_gradTrue) y x * 2 z y.detach() # z是一个新叶子requires_gradFalse w z * 3 # w的运算不会记录梯度因为z不需要梯度 loss y.sum() # loss的计算图包含 x - y loss.backward() # x.grad 2 # z.grad 和 w.grad 不存在因为它们不在计算图中与requires_grad_(False)的区别x.requires_grad_(False)是就地修改x本身影响之后所有以x为起点的运算。y x.detach()是创建了一个新的张量yx本身的性质不变y与x之后的计算图无关。detach()常用于将中间变量从计算图中取出用于可视化、评估或作为不需要梯度的输入传递给另一个网络而原始计算图的其他部分保持不变。3.2 torch.no_grad()上下文管理器临时禁用梯度追踪torch.no_grad()是一个上下文管理器。在其作用域内进行的所有运算无论输入张量的requires_grad是什么都不会被记录到计算图中。它主要用于推理阶段可以显著减少内存消耗因为不需要为反向传播存储中间变量。x torch.tensor([1.0], requires_gradTrue) with torch.no_grad(): y x * 2 # y.requires_grad False 且运算不被记录 # 即使x需要梯度这里的乘法也不会构建计算图 z x * 3 # 离开no_grad上下文这里会正常构建计算图与requires_grad_(False)的区别requires_grad_(False)是针对特定张量的持久性设置。torch.no_grad()是针对一段代码块的临时性设置影响其中所有张量的行为。它更轻量常用于模型验证 (model.eval()常与torch.no_grad()搭配使用)。3.3 使用场景对比表格操作作用对象是否就地主要用途对计算图的影响tensor.requires_grad_(False)单个张量是永久冻结模型参数如微调时冻结特征提取器。该张量之后参与的前向运算默认不构建梯度计算路径。tensor.requires_grad_(True)单个张量是解冻之前冻结的参数。仅影响该张量之后参与的前向运算。需同步更新优化器。tensor.detach()单个张量否返回新张量获取一个与当前张量数据相同但无梯度、无历史的新张量。常用于阻止梯度流向某一部分如GAN中固定生成器训练判别器。完全切断返回的新张量与之前计算图的联系。原张量不受影响。torch.no_grad()上下文内的所有运算不适用模型推理、计算评估指标、手动更新参数避免在优化器外产生梯度。临时禁用整个上下文内的梯度计算图构建。注意一个常见的误区是认为设置了requires_gradFalse或使用了no_grad()后计算速度会大幅提升。实际上主要节省的是内存因为不用存储中间变量的梯度信息和反向传播的计算开销。前向计算本身的开费几乎不变。4. 实战中的典型场景与避坑指南理解了原理我们来看看requires_grad_(True)在真实项目中的应用场景和那些容易踩的坑。4.1 场景一渐进式解冻与差分学习率在迁移学习中渐进式解冻是一种有效策略先冻结所有层训练分类头然后从顶层开始逐层解冻并训练。这里就涉及到requires_grad_(True)的批量操作。# 假设我们有一个预训练模型 pretrained_model # 第一步冻结所有层 for param in pretrained_model.parameters(): param.requires_grad False # 注意这里直接赋值属性也可以与调用方法等效。 # 第二步替换并训练新的分类头 classifier nn.Linear(pretrained_model.feature_dim, num_classes) model nn.Sequential(pretrained_model, classifier) # 此时只有classifier的参数需要梯度 # 训练几个epoch后解冻最后两个骨干网络块 unfreeze_layers [pretrained_model.layer4, pretrained_model.layer3] for layer in unfreeze_layers: for param in layer.parameters(): param.requires_grad_(True) # 解冻 # *** 关键步骤更新优化器 *** # 方法A重新创建优化器简单可靠 optimizer torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3) # 方法B如果使用像Adam这样有状态的优化器且想保留其他参数的状态操作会复杂很多。 # 通常更推荐方法A除非有特别理由需要保留优化器状态。避坑点优化器状态如上所述解冻后必须更新优化器。对于 Adam 等带有动量momentum和自适应学习率如exp_avg,exp_avg_sq的优化器新解冻的参数没有历史状态。如果强行将它们加入旧的优化器参数组可能会导致训练初期不稳定。新建优化器是最稳妥的做法。学习率设置解冻的层通常需要更小的学习率以防止破坏预训练的特征。这就是差分学习率。在新建优化器时可以通过传递一个字典列表给param_groups来实现。optimizer torch.optim.Adam([ {params: classifier.parameters(), lr: 1e-3}, # 分类头用较大学习率 {params: pretrained_model.layer3.parameters(), lr: 1e-4}, # 浅层用较小学习率 {params: pretrained_model.layer4.parameters(), lr: 1e-4}, ])4.2 场景二梯度检查与模型调试有时我们需要检查模型中特定部分的梯度是否正常传播。requires_grad_()可以用于临时“激活”某个中间层的梯度检查而不影响其他部分。# 假设我们怀疑某个中间层 mid_layer 的输出没有梯度 # 可以先将其输入和该层本身的参数梯度需求打开 original_requires_grad [] for param in mid_layer.parameters(): original_requires_grad.append(param.requires_grad) param.requires_grad_(True) # 执行一次前向和反向 output model(input) loss output.sum() loss.backward() # 检查 mid_layer 输出的梯度 if mid_layer.output is not None: print(mid_layer.output.grad) # 查看梯度是否存在且不为零 # 恢复原来的 requires_grad 状态 for param, original_flag in zip(mid_layer.parameters(), original_requires_grad): param.requires_grad_(original_flag)避坑点状态污染调试完毕后务必恢复参数的原始requires_grad状态否则会影响后续的正常训练流程。最好在修改前进行备份。计算图残留即使恢复了requires_grad这次调试运行构建的计算图可能已经增加了内存占用。在不需要时最好在torch.no_grad()上下文中运行此类检查或者直接使用loss.backward(retain_graphFalse)让计算图在反向传播后自动释放。4.3 场景三自定义操作与 in-place 操作的风险当你需要实现自定义的 PyTorch 函数通过继承torch.autograd.Function时或者进行原地in-place操作时对requires_grad的处理需要格外小心。原地操作的风险像tensor.add_(x)这样的原地操作如果tensor的requires_gradTrue并且x也需要梯度那么这个操作可能会破坏计算图导致梯度错误或RuntimeError。Autograd 很难追踪原地修改。一个安全的模式是在需要梯度计算的张量上尽量避免原地操作除非你非常清楚其后果。自定义 Function在编写forward和backward方法时你需要根据输入张量的requires_grad属性来决定是否需要计算和返回梯度。通常backward的输入是上游梯度你需要计算并返回本层所有需要梯度的输入的梯度。如果某个输入在forward时requires_gradFalse那么在backward中对应位置应返回None。class MyCustomFunction(torch.autograd.Function): staticmethod def forward(ctx, input): ctx.save_for_backward(input) # 保存供backward用的数据 # ... 执行前向计算 return output staticmethod def backward(ctx, grad_output): input, ctx.saved_tensors # 计算梯度 grad_input grad_output * some_operation(input) # 如果 forward 有多个输入这里需要返回对应数量的梯度。 # 对于不需要梯度的输入返回 None。 return grad_input # 使用 x torch.randn(5, requires_gradTrue) y MyCustomFunction.apply(x) # y的requires_grad会自动继承自x如果forward中所有输入都不需要梯度则y也不需要避坑点在自定义Function的backward中返回的梯度数量必须与forward的输入数量严格一致。对于不需要梯度的输入返回None。原地操作在自定义Function中尤其危险除非你明确知道自己在做什么并且处理好了梯度传播。5. 高级话题requires_grad 与模型部署、量化及 ONNX 导出在模型训练完成后requires_grad的角色并未结束。在部署和转换阶段它同样扮演着重要角色。5.1 推理优化统一设置 requires_gradFalse在将模型部署到生产环境时为了获得最佳性能和最小内存占用一个标准的做法是将整个模型设置为推理模式model.eval() # 将模型设置为评估模式影响Dropout、BatchNorm等层的行为 with torch.no_grad(): # 上下文管理器禁用梯度计算 for input in data_loader: output model(input) # ... 后续处理但torch.no_grad()是上下文管理器。如果你希望模型对象本身在全局状态下都不计算梯度例如将模型保存后加载到另一个完全不需要训练的环境中可以在导出前将所有参数的requires_grad设为Falsefor param in model.parameters(): param.requires_grad_(False) # 然后保存模型 torch.save(model.state_dict(), inference_model.pth)这样即使加载模型后不小心在非no_grad上下文中运行也不会产生多余的计算图和内存开销。不过更常见的做法是配合torch.jit.script或torch.jit.trace进行模型编译编译过程会自动处理梯度需求。5.2 模型量化与 requires_gradPyTorch 的动态量化Dynamic Quantization或静态量化Static Quantization通常用于训练后以减少模型大小和加速推理。量化过程本身如寻找缩放比例和零点可能需要微调Quantization-Aware Training, QAT。在 QAT 中你会在模型中插入伪量化FakeQuantize模块来模拟量化误差。这些伪量化模块通常需要梯度来进行微调。因此在准备 QAT 时你需要确保模型处于训练模式 (model.train())并且相关参数的requires_grad为True。量化工具链如torch.quantization.prepare_qat会帮你处理大部分设置。避坑点如果你在量化后加载模型进行纯推理记得将模型设置为eval()模式并且使用torch.no_grad()。量化后的模型其参数本身可能已经是requires_gradFalse的例如torch.quantization.convert后的模型但养成好习惯总是没错的。5.3 导出 ONNX 模型将 PyTorch 模型导出为 ONNX 格式时requires_grad的状态会影响导出结果。训练模式 vs 推理模式ONNX 导出通常期望模型处于推理模式 (model.eval())。在训练模式下一些具有随机性的操作如 Dropout会被导出这可能不是你想要在推理时看到的。梯度信息ONNX 图主要描述前向传播。即使你的 PyTorch 模型参数requires_gradTrue导出的 ONNX 模型默认也不包含反向传播信息。ONNX 可以导出用于梯度计算的图需要指定trainingTrainingMode.TRAINING或启用export_paramsTrue等但这主要用于支持某些需要梯度的后端或框架如 ONNX Runtime 的培训 API并非典型推理场景。一个稳健的 ONNX 导出流程是import torch model.eval() # 设置为评估模式 dummy_input torch.randn(1, 3, 224, 224) # 示例输入 # 在 no_grad 上下文中导出确保不构建计算图 with torch.no_grad(): torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, # 动态轴示例 )避坑点确保你的dummy_input的requires_grad属性与推理时的输入一致通常为False。如果模型包含条件判断如if语句依赖于输入数据导出可能会失败或产生不符合预期的静态图。对于复杂控制流可能需要使用torch.jit.script先编译再导出。6. 性能考量与最佳实践总结最后我们来聊聊requires_grad对性能的实际影响并总结一些黄金法则。6.1 内存与计算开销内存开销这是最主要的影响。对于requires_gradTrue的张量在前向传播中PyTorch 会为几乎所有中间计算结果非叶子节点保留一份副本用于反向传播时的梯度计算。这被称为“计算图保留”。对于一个深度网络这可能会使内存占用增加数倍。使用torch.no_grad()或设置requires_gradFalse可以避免这部分开销。计算开销反向传播本身需要计算梯度这当然会增加计算时间。但对于前向传播requires_grad的状态对计算速度影响微乎其微因为前向计算的核心运算是相同的。性能建议推理时务必使用model.eval()和torch.no_grad()这是最重要的性能优化习惯可以立即释放大量显存。及时释放不需要的计算图在训练循环中如果某些中间变量在后续不再需要可以考虑使用.detach()将其从计算图中分离或者确保loss.backward()调用后计算图被释放默认retain_graphFalse。梯度累加当 GPU 内存不足时一种技巧是使用梯度累加Gradient Accumulation。即在多个小批次micro-batch上计算损失和梯度但不立即调用optimizer.step()而是让梯度在.grad属性中累加。在累积了足够步数后再调用step()和optimizer.zero_grad()。在这个过程中所有参数的requires_grad必须为True。6.2 最佳实践清单基于以上的讨论我总结出以下处理requires_grad和requires_grad_()的最佳实践明确意图在修改requires_grad前想清楚你的目的是永久冻结/解冻参数还是临时禁用梯度计算。前者用requires_grad_()后者用torch.no_grad()。冻结/解冻后更新优化器这是最容易被忽略的坑。修改了参数的requires_grad属性后特别是解冻操作一定要重新创建优化器或者确保优化器的参数组包含了所有当前需要梯度的参数。对于简单的 SGD可能影响不大但对于 Adam 等复杂优化器这是必须的。使用model.requires_grad_()PyTorch 的nn.Module也有requires_grad_()方法可以递归地设置其所有子模块参数的requires_grad属性。这比手动遍历参数更简洁。# 冻结整个特征提取器 backbone.requires_grad_(False) # 只解冻分类层 classifier.requires_grad_(True)推理优先任何不需要训练/微调的代码路径如验证、测试、部署都放在model.eval()和torch.no_grad()上下文管理器中。谨慎使用原地操作在需要梯度的张量上尽量避免add_(),mul_()等原地操作除非你完全理解其对 Autograd 的影响。非原地操作如add()总是更安全。调试工具利用torch.autograd.gradcheck来检查自定义autograd.Function的梯度实现是否正确。在复杂模型梯度异常时可以使用torch.autograd.detect_anomaly()上下文管理器来帮助定位产生NaN或Inf梯度的操作。回到我最初遇到的问题根本原因就是忽略了“解冻参数需同步更新优化器”这一条。PyTorch 的自动微分系统设计得非常灵活和动态requires_grad是这个动态系统的阀门之一。理解它何时、如何影响计算图的构建与销毁而不仅仅是把它看成一个布尔标志能让你在模型训练、调试和部署中更加游刃有余避免很多隐晦的 bug。记住在 PyTorch 的世界里梯度流不是理所当然的它是由你通过requires_grad、计算图上下文和优化器共同精心控制的。