PyTorch Autograd 深入解析：从动态计算图到反向传播

摘要: Autograd 是 PyTorch 一切训练循环的地基。本文不绕弯子，直接用可运行的实验代码逐段拆解：requires_grad 如何触发建图、backward() 执行时图里发生了什么、梯度如何累积、detach()/no_grad() 何时必须用，最后落到显存与混合精度下的工程启示。

0. 实验目标

• 理解 PyTorch 如何把一次前向变成一张动态计算图（Dynamic Graph）
• 看清 backward() 沿图反向传播时，x.grad / w.grad 是怎么算出来的
• 用梯度累积实验验证"backward 不清零"的行为
• 明确 detach()、torch.no_grad()、inplace 操作的边界
• 给出训练工程（显存、AMP、梯度裁剪）中与 Autograd 直接相关的几条铁律

1. 第一段代码：一张最小的计算图

import torch

# 叶子张量：开启 requires_grad 才会被建图追踪
x = torch.tensor([2.0, 3.0], requires_grad=True)
w = torch.tensor([4.0, 5.0], requires_grad=True)

# 前向：y = x·w 逐元素相乘后求和
z = (x * w).sum()

print("z =", z.item())          # 2*4 + 3*5 = 23
print("z.requires_grad =", z.requires_grad)   # True

逐段拆解：

• requires_grad=True 之后，该张量成为计算图的叶子节点（leaf）。叶子节点是反向传播的"终点"，只有叶子节点的 .grad 会被填充。
• x * w 生成中间张量（非叶子），.sum() 生成 z。这三个算子执行时，PyTorch 同时把 x → mul → sum 的依赖关系记进了一张有向无环图（DAG）。
• 图里每个节点保存两类信息：前向输出值（反向要用）和反向函数（MulBackward / SumBackward）。
• z.requires_grad 为 True，因为它的输入链上存在需要梯度的张量——梯度标记会沿图传染。

2. 第二段代码：backward 到底做了什么

z.backward()

print("x.grad =", x.grad)   # tensor([4., 5.])
print("w.grad =", w.grad)   # tensor([2., 3.])

原理讲解：

z = Σ xᵢ·wᵢ，所以 ∂z/∂xᵢ = wᵢ，∂z/∂wᵢ = xᵢ。

backward() 从 z 出发，按 DAG 的逆拓扑序逐节点执行反向函数：

1. SumBackward：把上游梯度 1.0 广播到每个加数（[1, 1]）。
2. MulBackward：对 x 方向，上游梯度 × w → [1·4, 1·5] = [4, 5]，写入 x.grad；对 w 方向，上游梯度 × x → [2, 3]，写入 w.grad。
3. 叶子节点的 grad 填充完成后，图被释放（retain_graph=False 默认行为）。

关键点：梯度是"数值链式法则"的乘积，不是解析求导。每个算子自带反向函数，PyTorch 只是把上游梯度逐个乘过去。这就是为什么自定义 torch.autograd.Function 只需要实现 forward 和 backward 两个函数。

3. 第三段代码：梯度累积（backward 不清零）

x = torch.tensor([2.0], requires_grad=True)
w = torch.tensor([3.0], requires_grad=True)

for i in range(3):
    z = (x * w).sum()
    z.backward()          # 默认梯度是"累加"进 .grad

print("x.grad =", x.grad)   # tensor([9.])，即 3+3+3
print("w.grad =", w.grad)   # tensor([6.])，即 2+2+2

实验结果：

• 连续三次 backward()，x.grad 从 3.0 累加到 9.0——Autograd 默认不清零梯度。
• 这就是每个训练循环开头都要 optimizer.zero_grad() 的原因：不手动清零，梯度会跨 batch 累积，更新方向被污染。
• 反过来这也是"梯度累积（gradient accumulation）"技巧的根基：想模拟更大 batch 时，故意不清零、多跑几个 micro-batch 再 optimizer.step()。

工程启示：

# 梯度累积模拟大 batch：accum_steps 个 micro-batch 后再更新
optimizer.zero_grad()
for micro_batch in loader:
    loss = model(micro_batch)
    loss = loss / accum_steps          # 归一化，否则等效 batch 变大
    loss.backward()
    if (step + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

4. 第四段代码：detach 与 no_grad 的边界

x = torch.tensor([2.0], requires_grad=True)

# 情况 A：直接使用 x 参与后续计算
z1 = (x * 3).sum()
z1.backward()
print("A. x.grad =", x.grad)          # tensor([3.])

x.grad = None

# 情况 B：detach 切断梯度流（前向值还在，梯度不再回传）
y = x.detach()
z2 = (y * 3).sum()
z2.backward()
print("B. x.grad =", x.grad)          # tensor([0.]) —— 梯度被切断

原理讲解：

• detach() 返回一个共享同一份数据、但 requires_grad=False 且不在计算图里的新张量。反向经过它时梯度直接"断路"。
• 典型用途：
  • 对比学习/目标网络：Teacher 分支的梯度不应回传到 Student 分支（如 MoCo、BYOL、RL 里的 target Q-network）。
  • 特征冻结：预训练 backbone 提取特征后 detach() 再进分类头，backbone 不参与训练。
• torch.no_grad() 是上下文级别的切断：块内所有新建张量都不建图、不追踪，requires_grad 全部失效。推理（model.eval() + 预测）必须用它，省显存、省时间：

with torch.no_grad():
    pred = model(x)          # 不建图，不占中间激活显存

⚠️ inplace 操作的坑：对需要梯度的张量做 x.add_(1) 这类 inplace 修改，会破坏反向所需的旧值，PyTorch 会报 RuntimeError: a leaf Variable that requires grad is being used in an in-place operation。反向要用的是"前向时的旧值"，inplace 把旧值覆盖了，图就失效了。

5. 第五段代码：手动实现一个 backward 对比

import torch

class SquareAndDouble(torch.autograd.Function):
    """自定义算子：forward 输出 2·x²，backward 梯度 = 4·x·g"""
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return 2 * x * x

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * 4 * x

x = torch.tensor([3.0], requires_grad=True)
y = SquareAndDouble.apply(x)   # 2*9 = 18
y.backward()
print("y =", y.item(), "| x.grad =", x.grad.item())   # 4*3 = 12

实验结果：手写反向 12 = ∂(2x²)/∂x = 4x = 4×3，与链式法则一致。验证了"PyTorch 梯度 = 每个算子反向函数 × 上游梯度"的模型对自定义算子同样成立。

工程启示：算子融合、量化敏感算子、自定义 CUDA kernel 的梯度，都可以用 torch.autograd.Function 挂到 PyTorch 训练栈里，而不需要动训练框架本身。

6. 混合精度（AMP）下的 Autograd 行为

scaler = torch.cuda.amp.GradScaler()

with torch.autocast(device_type="cuda", dtype=torch.float16):
    loss = model(x)

scaler.scale(loss).backward()        # 先放大 loss，再 backward
scaler.step(optimizer)               # 放大后的梯度先 unscale 再更新
scaler.update()

原理讲解：

• AMP 下前向用 FP16 算（快、省显存），但梯度累积与参数更新仍建议 FP32，避免精度漂移。
• GradScaler 把 loss 放大（如 2¹⁶）后再 backward，是为了让 FP16 下容易下溢的小梯度在反传过程中保持可表示；step() 前自动 unscale，梯度里有 Inf/NaN 时跳过更新（scaler.update() 维护动态缩放因子）。
• 对 Autograd 而言，AMP 只是换了前向的 dtype 和反向的梯度缩放，建图与传播机制完全不变。

7. 工程启示汇总（可直接抄进团队规范）

1. 每个 step 前必 optimizer.zero_grad()：Autograd 默认累积梯度，不清零 = 梯度被污染。
2. 推理/评估用 torch.no_grad()：不建图，省一半以上激活显存，速度快 30%-50%。
3. 跨分支传值用 detach()：目标网络、对比学习、特征冻结，梯度切断处必须显式 detach。
4. 梯度累积模拟大 batch：多 micro-batch 不清零 + 归一化 loss + 按步更新。
5. AMP 必配 GradScaler：小梯度防下溢；scaler.step 前不要手动 unscale。
6. 自定义算子走 torch.autograd.Function：前向/反向成对实现，ctx.save_for_backward 保存中间值。
7. inplace 操作远离叶子张量：需要梯度的张量别 add_/mul_，报错即查图破坏。

8. 延伸阅读

• PyTorch 官方 Autograd 文档
• PyTorch 自定义 Function 教程
• 本知识库：09. 深度学习框架 章节