PyTorch计算图与自动求导的配合细节

PyTorch 的核心设计之一，就是把计算图（Computation Graph）和自动求导（Autograd）深度绑定在一起。

很多教程只讲：

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
y.backward()
print(x.grad)

得到：

4

但真正理解 PyTorch，需要知道：

PyTorch 不是“先建图再求导”，而是“运算过程中动态构建求导图，然后反向遍历图”。

下面深入到内部机制。

⸻

1. 什么是计算图

考虑：

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

a = x * 3
b = a + 4
y = b ** 2

数学表达：

y=(3x+4)^2

PyTorch 在执行过程中生成图：

x
|
Mul(3)
|
a
|
Add(4)
|
b
|
Pow(2)
|
y

更准确一点：

Tensor(x)
|
MulBackward0
|
Tensor(a)
|
AddBackward0
|
Tensor(b)
|
PowBackward0
|
Tensor(y)

这里已经出现一个关键概念：

图中的节点其实不是 Tensor，而是 backward Function。

⸻

2. Tensor 与 Function 的关系

查看：

print(y.grad_fn)

输出类似：

<PowBackward0>

查看：

print(b.grad_fn)

输出：

<AddBackward0>

查看：

print(a.grad_fn)

输出：

<MulBackward0>

而：

print(x.grad_fn)

输出：

None

因为：

x

是叶子节点（Leaf Tensor）。

⸻

内部关系：

Tensor(y)
|
grad_fn
|
PowBackward0

而：

PowBackward0
|
next_functions
|
AddBackward0

继续：

AddBackward0
|
MulBackward0
|
AccumulateGrad(x)

⸻

3. next_functions

可以直接看：

print(y.grad_fn.next_functions)

得到：

((<AddBackward0 object>, 0),)

继续：

print(y.grad_fn.next_functions[0][0].next_functions)

得到：

((<MulBackward0 object>, 0),)

再继续：

mul = y.grad_fn.next_functions[0][0].next_functions[0][0]

print(mul.next_functions)

会看到：

((<AccumulateGrad object>, 0),)

说明：

PowBackward0
↓
AddBackward0
↓
MulBackward0
↓
AccumulateGrad
↓
x.grad

反向传播本质上就是沿着这个链表向后走。

⸻

4. 动态计算图（Dynamic Graph）

PyTorch 与早期 TensorFlow 最大区别：

TensorFlow 1.x：先定义图,再运行图;

PyTorch：运行即建图;

例如：

for i in range(3):
y = x ** i

每次循环：

重新创建图
重新销毁图

因此：

if condition:
y = x * 2
else:
y = x * 3

也是天然支持的。

因为图是运行时生成。

⸻

5. backward() 到底做了什么

假设：

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

y = (3*x+4)**2

图：

x
│
Mul
│
Add
│
Pow
│
y

调用：

y.backward()

实际上：

Autograd Engine.run_backward()

启动。

流程：

第一步

初始化：

\frac{\partial y}{\partial y}=1

即：

grad_output = 1

⸻

第二步

进入 PowBackward

对于：

z=b^2

有：

\frac{\partial z}{\partial b}=2b

当前：

b = 10

所以：

2b=20

得到：

\frac{\partial y}{\partial b}=20

⸻

第三步

进入 AddBackward

对于：

b=a+4

有：

\frac{\partial b}{\partial a}=1

所以：

\frac{\partial y}{\partial a}
=
20\times1
=
20

⸻

第四步

进入 MulBackward

对于：

a=3x

有：

\frac{\partial a}{\partial x}=3

因此：

\frac{\partial y}{\partial x}
=
20\times3
=
60

⸻

第五步

AccumulateGrad

写入：

x.grad = 60

结束。

⸻

6. 为什么需要保存中间值

考虑：

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

y = x**3

反向：

\frac{dy}{dx}=3x^2

Backward 时必须知道：

x=2

因此 Forward 时：

PowBackward0

会保存：

saved_tensors

内部类似：

ctx.save_for_backward(x)

因此：

y.grad_fn

实际上带着 Forward 的部分数据。

⸻

7. Saved Tensor 机制

自定义算子：

class Cube(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return x**3

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

这里：

ctx

就是连接：

Forward
↓
保存中间值
↓
Backward

的桥梁。

⸻

8. 梯度累积机制

很多人第一次会疑惑：

y.backward()
y.backward()

结果：

x.grad = 120

而不是：

60

因为：

x.grad += new_grad

不是覆盖。

内部：

AccumulateGrad

节点负责：

grad += incoming_grad

所以训练前必须：

optimizer.zero_grad()

或：

model.zero_grad()

⸻

9. 为什么默认只能 backward 一次

y.backward()
y.backward()

报错：

Trying to backward through the graph a second time

原因：

第一次 backward 后：

saved tensors
释放
graph释放

节省显存。

⸻

如果想保留：

y.backward(retain_graph=True)

图不释放。

⸻

10. 叶子节点（Leaf Tensor）

x = torch.tensor(1., requires_grad=True)

y = x * 2

这里：

x.is_leaf

True

y.is_leaf

False

⸻

只有叶子节点默认保存梯度：

print(x.grad)

有值

print(y.grad)

None

⸻

如果需要：

y.retain_grad()

之后：

y.backward()

即可获得：

y.grad

⸻

11. detach 的本质

y = x * 2

z = y.detach()

变成：

x ──► y

z

而不是：

x ──► y ──► z

即：

z.requires_grad

为 False。

反向传播到：

z

不会回到：

x

⸻

12. no_grad 的本质

with torch.no_grad():
y = x * 2

此时：

y.grad_fn

为：

None

因为：

Autograd Engine
停止记录运算

图根本没建。

⸻

13. View 与计算图

x = torch.randn(4, requires_grad=True)

y = x.view(2,2)

这里：

y

与

x

共享存储。

图：

x
│
ViewBackward
│
y

反向时：

y.grad

最终会 reshape 回：

x.grad

原始形状。

⸻

14. Autograd Engine 的真实执行模型

实际并非递归。

PyTorch C++ 中采用：

GraphTask
ReadyQueue
Worker Threads

结构。

反向传播时：

Root Node
↓
拓扑排序
↓
Ready Queue
↓
多线程执行Backward Node

类似：

 A
/  \
B  C
\  /
 D

B、C 的梯度计算可以并行。

这也是大型模型反向传播效率高的重要原因。

⸻

15. 整个生命周期总结

执行：

x = torch.tensor(2., requires_grad=True)

y = ((x*3)+4)**2

发生：

Forward

Tensor(x)

↓
MulBackward0

↓
AddBackward0

↓
PowBackward0

↓
Tensor(y)

同时：

保存反向所需中间值

⸻

Backward

dy/dy = 1

↓
PowBackward0

↓
AddBackward0

↓
MulBackward0

↓
AccumulateGrad

应用链式法则：

\frac{dy}{dx}
=
\frac{dy}{db}
\frac{db}{da}
\frac{da}{dx}

最终：

x.grad = 60

⸻

从源码视角看，PyTorch 自动求导系统可以概括为：

Tensor
↓
记录运算
↓
构建 Function DAG
↓
Forward 保存必要状态
↓
backward() 启动 Engine
↓
逆拓扑遍历 DAG
↓
链式法则传播梯度
↓
AccumulateGrad 写入叶子节点

理解这一套之后，再去看 PyTorch 源码中的 TensorImpl、AutogradMeta、Node(Function)、Engine::execute()、自定义 torch.autograd.Function，就会发现它们其实是在围绕同一个核心：动态构图 + 反向模式自动微分（reverse-mode autodiff）。