04_torch_no_grad_vs_require_grad=False
本节主要探索这样一个问题:Pytorch 中,torch.no_grad() 与 param.requires_grad==False 有什么区别?
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 = FalseWhen 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 Successfullyb. Yes No 112M Successfullyc. No Yes 3M Successfullyd. No No 112M CUDA out of memoryCan 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.grad为None) - ✅ 优化器(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 仍为 None1.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_fn1.3 为什么组合 (d) 会 OOM,但 (b) 不会?
回到原帖主提到的四个组合:
| 组合 | no_grad | requires_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 在背后偷偷建了这样一张图:
- 每个运算都会创建一个
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 = 5z = 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.0print(b.grad) # ∂z/∂b = c = 5.0print(c.grad) # ∂z/∂c = a + b = 5.02.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 = 52.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「这笔账不用记」,从而省下记帐的内存。
文章分享
如果这篇文章对你有帮助,欢迎分享给更多人!