04_torch_no_grad_vs_require_grad=False

1746 字
9 分钟
04_torch_no_grad_vs_require_grad=False

本节主要探索这样一个问题:Pytorch 中,torch.no_grad()param.requires_grad==False 有什么区别

原问题来源 Stack Overflow


I’m following a PyTorch tutorial which uses the BERT NLP model (feature extractor) from the Huggingface Transformers library. There are two pieces of interrelated code for gradient updates that I don’t understand.

(1) torch.no_grad()

The tutorial has a class where the forward() function creates a torch.no_grad() block around a call to the BERT feature extractor, like this:

bert = BertModel.from_pretrained('bert-base-uncased')
class BERTGRUSentiment(nn.Module):
def __init__(self, bert):
super().__init__()
self.bert = bert
def forward(self, text):
with torch.no_grad():
embedded = self.bert(text)[0]

(2) param.requires_grad = False

There is another portion in the same tutorial where the BERT parameters are frozen.

for name, param in model.named_parameters():
if name.startswith('bert'):
param.requires_grad = False

When would I need (1) and/or (2)?

  • If I want to train with a frozen BERT, would I need to enable both?
  • If I want to train to let BERT be updated, would I need to disable both?

Additionaly, I ran all four combinations and found:

with torch.no_grad requires_grad = False Parameters Ran
------------------ --------------------- ---------- ---
a. Yes Yes 3M Successfully
b. Yes No 112M Successfully
c. No Yes 3M Successfully
d. No No 112M CUDA out of memory

Can someone please explain what’s going on? Why am I getting CUDA out of memory for (d) but not (b)? Both have 112M learnable parameters.


这是一个非常经典的问题。那四个组合的实验结果本身就很好地说明了区别,我们来逐一拆解:


1.1 本质区别#

概念作用范围本质
param.requires_grad = False单个参数/张量告诉 autograd:“这个参数我不需要梯度”
torch.no_grad()整个代码块告诉 autograd:“这段代码里所有操作都不建计算图”

1.2 详细解释#

1.2.1 requires_grad = False —— 参数级别的开关#

这是某个张量的属性,表示 PyTorch 不需要为该张量计算梯度。

  • ✅ 该参数本身的梯度不会被计算(param.gradNone
  • ✅ 优化器(SGD、Adam 等)不会更新这个参数
  • 但经过该参数的计算仍然会构建计算图,只要下游还有需要梯度的张量
  • 仍然会消耗内存存储中间激活值,以备反传使用
x = torch.randn(3, requires_grad=True)
w = torch.randn(3, requires_grad=False) # 这个参数不需要梯度
y = (w * x).sum() # 计算图仍然会构建!
y.backward() # x.grad 被计算,w.grad 仍为 None

1.2.2 torch.no_grad() —— 全局上下文管理器#

进入这个上下文后,所有操作都不再追踪梯度,完全不构建计算图。

  • 彻底关闭计算图构建,不存储任何中间激活值
  • ✅ 大幅节省 GPU 显存(backward 需要的内存 ≈ forward 的 2~3 倍)
  • ✅ 无论张量的 requires_grad 是什么,都不计算梯度
  • ❌ 如果你在 no_grad() 里修改了参数,优化器后续更新时也不会感知这些修改
x = torch.randn(3, requires_grad=True)
w = torch.randn(3, requires_grad=True)
with torch.no_grad():
y = (w * x).sum() # 完全不建计算图!
# y.backward() ❌ 会报错,因为 y 没有 grad_fn

1.3 为什么组合 (d) 会 OOM,但 (b) 不会?#

回到原帖主提到的四个组合:

组合no_gradrequires_grad参数量结果
a✅ (False)3M
b❌ (True)112M
c✅ (False)3M
d❌ (True)112M❌ OOM
╔══════════════════════════════════════════════════════════════╗
║ (b) torch.no_grad() 激活 ║
║ ┌─────────────────────────────────────────────────────────┐ ║
║ │ BERT 前向传播 → 输出嵌入 │ ║
║ │ (不建计算图,不存中间激活,省显存) │ ║
║ └─────────────────────────────────────────────────────────┘ ║
║ GRU + 分类头 → 计算图(只有 3M 参数需反传) ║
║ 总显存 ≈ BERT 推理 + 3M 参数训练 ≈ 低 ║
╚══════════════════════════════════════════════════════════════╝
╔══════════════════════════════════════════════════════════════╗
║ (d) 没有 torch.no_grad() ║
║ ┌─────────────────────────────────────────────────────────┐ ║
║ │ BERT 前向传播 │ ║
║ │ (建计算图!12层 Transformer 的所有中间激活全部存入显存) │ ║
║ │ 需要保留:每层 Q/K/V/注意力分数/FFN 中间值... │ ║
║ └─────────────────────────────────────────────────────────┘ ║
║ 总显存 ≈ BERT 训练(正向 + 反向中间值)≈ 12GB+ → OOM ║
╚══════════════════════════════════════════════════════════════╝

关键点requires_grad=True 只是说”参数需要梯度”,但光靠它本身并不会让显存爆炸。真正让显存爆炸的是没有 no_grad() 时,PyTorch 会为 BERT 的每一层都保存中间激活值以便反传——12 层 Transformer 的中间结果是非常吃显存的。


1.4 实际场景总结#

场景需要 no_grad需要 requires_grad=False原因
冻结 BERT 做特征提取✅ 建议用✅ 建议用no_grad 省显存,requires_grad=False 防止意外更新 + 优化器跳过
微调 BERT❌ 不能用❌ 不能用需要计算图和梯度来更新 BERT
仅推理/评估✅ 必须用❌ 不需要no_grad 省显存提速,参数是否可训练无关紧要

结论requires_grad=False 控制是否更新参数(优化器角度),torch.no_grad() 控制是否建计算图(显存/速度角度)。在冻结 BERT 时,两者都用是最佳实践——就像教程中做的那样。

2. 什么是计算图#

计算图(Computation Graph)是 PyTorch 等深度学习框架的核心概念。它是 autograd(自动求导)的底层数据结构


2.1 直观理解#

计算图就是一个有向无环图(DAG),记录了你的每一步数学运算:

  • 节点(Node) = 张量(数据)或运算(函数)
  • 边(Edge) = 数据流向

当执行 z = (a + b) × c 时,PyTorch 在背后偷偷建了这样一张图:

graph TD subgraph "前向传播(从左到右建图)" A["a<br/>requires_grad=True"] --> Add["➕ Add<br/>grad_fn=<AddBackward0>"] B["b<br/>requires_grad=True"] --> Add Add --> MUL["✖️ Mul<br/>grad_fn=<MulBackward0>"] C["c<br/>requires_grad=True"] --> MUL MUL --> Z["z = (a+b)×c"] end
  • 每个运算都会创建一个 grad_fn(梯度函数),记录如何计算该运算的偏导
  • 这张图就是 链式法则 的地图 —— 反向传播时 PyTorch 按图索骥,从 z 一路往回走,依次调用每个 grad_fn 算出梯度

2.2 用一个具体例子演示#

import torch
a = torch.tensor([2.0], requires_grad=True)
b = torch.tensor([3.0], requires_grad=True)
c = torch.tensor([5.0], requires_grad=True)
# 前向传播:PyTorch 自动构建计算图
t = a + b # t = 2 + 3 = 5
z = t * c # z = 5 * 5 = 25
# 查看 grad_fn,看到计算图的"足迹"
print(t.grad_fn) # <AddBackward0 object>
print(z.grad_fn) # <MulBackward0 object>
# 反向传播:沿着计算图走,计算每个参数的梯度
z.backward()
print(a.grad) # ∂z/∂a = c = 5.0
print(b.grad) # ∂z/∂b = c = 5.0
print(c.grad) # ∂z/∂c = a + b = 5.0

2.2.1 反向传播的路线图#

z.backward() 开始
MulBackward0: 计算 ∂z/∂t = c, ∂z/∂c = t
AddBackward0: 计算 ∂z/∂a = ∂z/∂t × ∂t/∂a = c × 1 = 5
计算 ∂z/∂b = ∂z/∂t × ∂t/∂b = c × 1 = 5

2.3 计算图与内存#

这就是为什么之前说「建计算图很吃显存」—— 每个 grad_fn 必须记住它前向时的输入值(即中间激活),才能在反向时计算出梯度。

回到 BERT 的例子:

BERT 一层 Transformer 的计算图(简化):
输入 → QKV 线性变换 → 注意力分数 → Softmax → 加权求和 → FFN → 输出
↓ ↓ ↓ ↓ ↓ ↓
存了 QKV 存了分数 存了概率 存了加权 存了中间 存了输出

12 层 × 每层多个中间值 × batch_size × 序列长度 × 768 维 = 天文数字的显存

torch.no_grad() 就是告诉 PyTorch:「这段代码我不需要反向传播,别建图,别存中间值」—— 于是前向跑完直接丢弃中间结果,显存瞬间降下来。


2.4 一句话总结#

计算图就是 PyTorch 的「记忆账本」—— 前向时一笔笔记录每步运算和中间结果,反向时按图索骥算出所有梯度。torch.no_grad() 就是告诉 PyTorch「这笔账不用记」,从而省下记帐的内存。

文章分享

如果这篇文章对你有帮助,欢迎分享给更多人!

04_torch_no_grad_vs_require_grad=False
https://github.com/chunhuizhang/bilibili_vlogs/
作者
HAC
发布于
2026-06-09
许可协议
CC BY-NC-SA 4.0

评论区

Profile Image of the Author
HAC
观之非易,行且克难
Greetings
欢迎来到我的博客!这里主要分享我的学习笔记与兴趣爱好。
音乐
封面

音乐

暂未播放

0:00 0:00
暂无歌词
分类
标签
站点统计
文章
32
分类
5
标签
13
总字数
79,889
运行时长
0
最后活动
0 天前

文章目录