第4章 梯度与计算图：计算图原理、梯度存储机制、retain_graph与detach的妙用
说到深度学习，绕不开的就是梯度。而梯度的背后，是计算图在默默工作。
我刚开始接触PyTorch时，对计算图的理解很模糊。只知道调backward()就能算梯度，至于它怎么算的、内存怎么管理的，完全没概念。直到有一次，我训练一个稍复杂的模型，显存直接爆了。排查半天，发现是计算图没释放。嗯，从那以后，我开始认真研究这玩意儿。
计算图：到底是个啥？
说白了，计算图就是记录你所有运算过程的一张“流程图”。
你想想看，我们做前向传播时，每个张量是怎么来的？
a = torch.tensor([2.0], requires_grad=True)— 这是叶子节点b = a * 3— 这是运算结果c = b ** 2— 继续运算
PyTorch会在背后偷偷记下：a → 乘法 → b → 平方 → c。这就是一张有向无环图（DAG）。
我习惯把计算图想象成“案发现场重建”。前向传播时，每个操作都留下线索。反向传播时，就顺着这些线索往回走，用链式法则算出每个节点的梯度。
核心要点：计算图只在requires_grad=True的张量参与运算时才会构建。如果你用torch.no_grad()包裹，就不会建图，自然也不会有梯度。
下面这张图，展示了计算图的基本结构：
梯度存储机制：显存里的隐形杀手
每次调用backward()，PyTorch都会计算梯度，然后默认把梯度累加到张量的.grad属性上。
但这里有个坑：计算图本身也会占用显存。
为什么？因为PyTorch需要保存中间结果，才能在反向传播时计算梯度。比如上面的例子，计算c对b的梯度需要b的值，计算b对a的梯度需要a的值。这些中间结果，都保存在计算图里。
我曾经踩过的坑：训练一个GAN，生成器和判别器交替训练。每次调用backward()后，计算图还留在内存里。迭代了几百步，显存从8GB涨到12GB，最后直接OOM。排查了半天，才发现是计算图没释放。
默认情况下，每次backward()执行完后，计算图会被释放（free）。但如果你需要多次反向传播，就得用retain_graph=True。
retain_graph：什么时候用？
retain_graph，顾名思义，就是保留计算图。
什么时候需要？最常见的就是GAN训练：
# 生成器前向
fake = generator(z)
# 判别器判断真假
d_real = discriminator(real)
d_fake = discriminator(fake.detach())
# 判别器损失
d_loss = d_real_loss(d_real, real_labels) + d_fake_loss(d_fake, fake_labels)
d_loss.backward() # 第一次反向传播
optimizer_d.step()
# 生成器损失
g_loss = g_loss(discriminator(fake), real_labels)
g_loss.backward() # 第二次反向传播，需要保留计算图
optimizer_g.step()
等等，上面这段代码其实有问题。第一次backward()后，计算图就被释放了。第二次backward()时，discriminator(fake)的计算图已经没了。
正确的做法：
d_loss.backward(retain_graph=True) # 保留计算图
optimizer_d.step()
g_loss.backward() # 这次可以正常释放了
optimizer_g.step()
我的建议：能用detach就别用retain_graph。retain_graph会让计算图一直留在内存里，显存压力很大。我一般只在必须多次反向传播时才用，用完马上释放。
detach：切断计算图的利器
detach()的作用很简单：从计算图中分离出一个张量。分离后的张量不再需要梯度，也不会参与反向传播。
我举个例子：
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
z = y.detach() # z不再需要梯度
w = z ** 2
w.backward() # 梯度只传到y，不会传到x
print(x.grad) # None
detach在强化学习中特别常用。比如PPO算法，旧策略采样的数据需要detach掉，不让梯度流回去。
还有一个经典场景：特征提取 + 分类器微调。
# 预训练的特征提取器
features = backbone(inputs)
# detach掉，只训练分类器
features = features.detach()
outputs = classifier(features)
loss = criterion(outputs, labels)
loss.backward()
这样，backward()时梯度只更新分类器，不会动backbone的参数。既省显存，又防止特征提取器被破坏。
核心区别总结：
| 方法 | 作用 | 显存影响 | 典型场景 |
|---|---|---|---|
| retain_graph=True | 保留计算图，允许多次backward | 增加显存占用 | GAN、多任务学习 |
| detach() | 切断张量与计算图的连接 | 减少显存占用 | 特征提取、强化学习 |
| torch.no_grad() | 上下文管理器，整个代码块不建图 | 大幅减少显存 | 推理、评估 |
实战避坑指南
我总结几个常见的坑，都是我在项目中真实遇到过的：
- 梯度累加时忘记清零：每次backward()后，梯度会累加到.grad上。如果你做梯度累加（gradient accumulation），记得手动zero_grad()。不然梯度会越加越大。
- in-place操作破坏计算图：比如x.add_(1)，这种操作会修改原始张量，导致计算图出错。PyTorch会报错：RuntimeError: a leaf Variable that requires grad has been used in an in-place operation.
- 循环中创建计算图不释放：如果你在循环里反复创建计算图，又不及时释放，显存会持续增长。我习惯在循环末尾显式调用del variable，或者用with torch.no_grad()包裹不需要梯度的部分。
- 多GPU时计算图复制：使用DataParallel时，每个GPU上都会复制一份计算图。显存占用直接翻倍。我建议用DistributedDataParallel，效率更高，显存也更省。
一个小技巧：调试时可以用torchviz库可视化计算图。安装后调用make_dot(loss)就能生成计算图PDF。我第一次看到自己模型的计算图时，才发现有些分支根本不需要梯度，白白浪费了显存。
嗯，关于计算图和梯度，今天就聊这么多。记住一句话：计算图是梯度的载体，也是显存的消耗者。用好retain_graph和detach，你就能在显存和功能之间找到平衡。
想要更多资料，联系可可