Table of Contents
5010 words
25 minutes
Autograd & Computational Graph —— 自动微分与计算图
Autograd & Computational Graph —— 自动微分与计算图
一、前言：AI框架最核心的“魔法”
在深度学习的日常开发中，我们几乎每天都在做同一件事：调用 loss.backward() 。这短短一行代码，框架就自动帮我们算好了所有参数的梯度，仿佛魔法一般。
但如果你曾好奇过：这个“魔法”到底是怎么实现的？为什么PyTorch能自动算出梯度？retain_graph和create_graph又是干什么的？PyTorch和TensorFlow在底层设计上到底有什么本质不同？
这些问题的答案，都指向同一个核心概念——自动微分（Automatic Differentiation, Autograd） 与计算图（Computational Graph） 。
本文将从五个方面，带你彻底搞懂自动微分与计算图的底层原理：
- 前向计算图与反向传播的链式法则映射
- 雅可比矩阵（Jacobian）与海森矩阵（Hessian）
- PyTorch中
retain_graph与create_graph的高阶导数应用场景 - 静态图（TF1.x）vs 动态图（PyTorch）的底层区别
- 实战：手写一个微型自动微分框架
二、计算图：自动微分的“骨架”
2.1 什么是计算图？
计算图（Computational Graph） 是一种有向无环图（DAG） ，用于表示一个计算过程。图中的节点（Node） 代表变量或操作（如加法、乘法、矩阵乘法），边（Edge） 代表数据流向。
1import torch2
3# 一个简单的计算：y = (x + 2) * 34x = torch.tensor([1.0], requires_grad=True)5a = x + 2 # 加法节点6y = a * 3 # 乘法节点7
8# 此时PyTorch内部已经构建了一个计算图：9# x(叶子节点) --(+2)--> a(中间节点) --(*3)--> y(根节点)任何复杂的神经网络，最终都可以表示为一个巨大的计算图——输入是数据，输出是损失值，中间是层层叠叠的矩阵乘法和非线性激活函数。
2.2 前向传播：沿着图“正向”计算
前向传播（Forward Propagation） 就是按照计算图的依赖顺序，从输入到输出依次计算每个节点的值。
以上面的例子为例：
- 输入
x = 1 - 计算
a = x + 2 = 3 - 计算
y = a * 3 = 9
在PyTorch中，前向传播的同时，框架会记录所有操作的顺序和中间结果。这些记录构成了反向传播的基础。注意，反向传播会重复利用前向传播中存储的中间值，这也是训练比推理需要更多显存的原因之一。
1# 查看计算图的结构2print(y.grad_fn) # <MulBackward0 object at 0x...>3print(a.grad_fn) # <AddBackward0 object at 0x...>4print(x.grad_fn) # None (叶子节点没有grad_fn)grad_fn 属性记录了创建该张量的操作，正是这些 grad_fn 链构成了完整的计算图。
2.3 反向传播：沿着图“逆向”求导
反向传播（Backward Propagation） 则是按照计算图的逆序，从输出向输入逐层计算梯度。
在训练神经网络时，前向传播和反向传播相互依赖。前向传播计算并存储中间变量，反向传播按相反的顺序（从输出层到输入层）计算这些中间变量和参数的梯度。
1# 反向传播：从y开始逆向求导2y.backward()3
4# 查看梯度5print(x.grad) # tensor([3.0])6# 数学验证：dy/dx = 3，因为 y = 3*(x+2) = 3x + 6三、链式法则：反向传播的“数学引擎”
3.1 链式法则在计算图上的应用
反向传播的数学基础就是微积分中的链式法则（Chain Rule） 。自动微分通过系统性地对计算图中的每个操作应用链式法则，来计算任意复杂函数的精确导数。
链式法则：如果 y=f(u) 且 u=g(x)，则：
dxdy=dudy⋅dxdu
在计算图中，一条路径的偏微分等于路径上各相邻节点之间偏微分的连乘。
3.2 多变量链式法则
在神经网络中，我们处理的是多变量函数。假设有计算图：
y=f(x1,x2,...,xn)
对于某个中间变量 z，其梯度传播遵循多变量链式法则：
∂xi∂y=∑k∂zk∂y⋅∂xi∂zk
这个求和对应了计算图中从 xi 到 y 的所有路径的贡献之和。
3.3 反向传播 vs 自动微分
反向传播（Backpropagation） 和自动微分（Automatic Differentiation, AD） 这两个概念经常被混用，但它们的侧重点略有不同：
- 反向传播：特指神经网络训练中从损失函数到各层参数的梯度计算过程
- 自动微分：是一个更通用的概念，指任何给定计算图后计算梯度的技术
在AI框架中，自动微分是实现反向传播的技术手段。自动微分将复合函数分解为计算图，并以此计算任意两个节点间的梯度。
自动微分有两种执行模式：
- 前向模式（Forward Mode） ：与函数结果一起计算梯度
- 反向模式（Reverse Mode） ：先评估函数，再从输出反向计算梯度
深度学习中使用的是反向模式，因为神经网络的输出是标量（损失值），输入是海量参数——反向模式的计算成本与输出维度成正比，而输出维度为1，效率远高于前向模式。
3.4 完整推导示例
用一个具体例子来演示完整的链式法则推导：
1计算图：y = (x₁ + x₂) * x₃2其中：x₁ = 2, x₂ = 3, x₃ = 41import torch2
3x1 = torch.tensor(2.0, requires_grad=True)4x2 = torch.tensor(3.0, requires_grad=True)5x3 = torch.tensor(4.0, requires_grad=True)6
7# 前向传播8a = x1 + x2 # a = 59y = a * x3 # y = 2010
11# 反向传播12y.backward()13
14print(f"dy/dx1 = {x1.grad.item()}") # 4.0 = x315print(f"dy/dx2 = {x2.grad.item()}") # 4.0 = x316print(f"dy/dx3 = {x3.grad.item()}") # 5.0 = a = x1 + x217
18# 数学验证：19# ∂y/∂x1 = ∂(a*x3)/∂x1 = x3 * ∂(x1+x2)/∂x1 = 4 * 1 = 4 ✓20# ∂y/∂x3 = ∂(a*x3)/∂x3 = a = x1 + x2 = 5 ✓四、雅可比矩阵与海森矩阵
4.1 雅可比矩阵（Jacobian Matrix）
当输出不是标量时，梯度就变成了一个矩阵——这就是雅可比矩阵（Jacobian Matrix） 。
对于函数 f:Rn→Rm，雅可比矩阵是一个 m×n 的矩阵：
\frac{\partial f_1}{\partial x_1} & \cdots & \frac{\partial f_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \frac{\partial f_m}{\partial x_1} & \cdots & \frac{\partial f_m}{\partial x_n} \end{bmatrix}$$ 其中 $J[i, j] = \partial f_i / \partial x_j$。 **直观理解**：雅可比矩阵描述了**输出向量每一个分量对输入向量每一个分量的敏感度**。 在深度学习中，PyTorch的Autograd实际计算的是**向量-雅可比积（Vector-Jacobian Product, VJP）** ，而不是完整的雅可比矩阵。这是因为： - 损失函数 $L$ 是标量，对输出 $\mathbf{y}$ 的梯度 $\nabla_{\mathbf{y}} L$ 是一个向量 - 我们真正需要的是 $\nabla_{\mathbf{x}} L = (\nabla_{\mathbf{y}} L)^T \cdot J_{\mathbf{f}}(\mathbf{x})$ - 即**向量与雅可比矩阵的乘积**，而非完整的雅可比矩阵 这大大提高了计算效率——我们不需要显式构建雅可比矩阵，只需要计算**向量-雅可比积**即可。 ```python import torch # 多输出函数：f(x) = [x₁² + x₂, x₁ - x₂²] def f(x): return torch.stack([x[0]**2 + x[1], x[0] - x[1]**2]) x = torch.tensor([2.0, 3.0], requires_grad=True) y = f(x) # y是向量（非标量），需要传入grad_tensors作为“向量” # 这里用 [1.0, 0.0] 表示只关注第一个输出的梯度 y.backward(gradient=torch.tensor([1.0, 0.0])) print(x.grad) # tensor([4.0, 1.0]) # 这是向量-雅可比积的结果：∂y₁/∂x = [2x₁, 1] = [4, 1] ``` ### 4.2 海森矩阵（Hessian Matrix） **海森矩阵（Hessian Matrix）** 是**标量函数对所有输入变量的二阶偏导数**组成的方阵： $$H = \begin{bmatrix} \frac{\partial^2 f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1 \partial x_2} & \cdots \\ \frac{\partial^2 f}{\partial x_2 \partial x_1} & \frac{\partial^2 f}{\partial x_2^2} & \cdots \\ \vdots & \vdots & \ddots \end{bmatrix}$$ 从另一个角度看，**海森矩阵是雅可比矩阵的雅可比矩阵**。 **海森矩阵的用途**： 1. **判断凸性**：如果海森矩阵是正定的，函数是凸的 2. **牛顿法优化**：使用海森矩阵的逆来加速收敛 3. **理解损失 landscapes**：海森矩阵的特征值反映了损失函数曲面的曲率 ```python # 计算海森矩阵（需要create_graph=True） x = torch.tensor([2.0, 3.0], requires_grad=True) y = x[0]**2 + x[1]**2 + x[0]*x[1] # 标量函数 # 第一步：计算梯度（保留计算图以便高阶求导） grad = torch.autograd.grad(y, x, create_graph=True)[0] # [2x₁+x₂, 2x₂+x₁] # 第二步：对梯度再求导得到海森矩阵 hessian = [] for i in range(2): hess_row = torch.autograd.grad(grad[i], x, retain_graph=True)[0] hessian.append(hess_row) hessian = torch.stack(hessian) print(hessian) # tensor([[2., 1.], # [1., 2.]]) # 验证：∂²y/∂x₁² = 2, ∂²y/∂x₁∂x₂ = 1, ∂²y/∂x₂² = 2 ``` 计算雅可比矩阵或海森矩阵在**非传统深度学习模型**（如物理信息神经网络PINN、隐式微分等）中非常有用。不过，使用标准的Autograd API高效计算这些量是困难且繁琐的，通常需要借助`functorch`等工具。 ## 五、PyTorch中的`retain_graph`与`create_graph` ### 5.1 计算图的默认行为：用完即焚 在PyTorch中，**默认情况下，一次 `backward()` 调用结束后，计算图会被立即释放（freed）** 。这是为了**节省内存**——计算图可能非常庞大，保留它会占用大量显存。 ```python x = torch.tensor([1.0], requires_grad=True) y = x ** 2 y.backward() # 计算图被释放 print(x.grad) # tensor([2.0]) # 再次调用backward会报错，因为图已经被释放 # y.backward() # RuntimeError: Trying to backward through the graph a second time... ``` ### 5.2 `retain_graph=True`：保留计算图供多次反向传播 **`retain_graph=True`** 告诉PyTorch：**反向传播结束后，不要释放计算图**。 **适用场景**： 1. **需要对同一个计算图进行多次反向传播**（如某些对抗训练算法） 2. **需要多次查看中间变量的梯度** 3. **累积梯度场景** ```python x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) y = x.pow(2).sum() # 第一次反向传播，保留计算图 y.backward(retain_graph=True) print(x.grad) # tensor([2.0, 4.0, 6.0]) # 第二次反向传播（梯度会累积！） y.backward(retain_graph=True) print(x.grad) # tensor([4.0, 8.0, 12.0]) ← 累加了！ # 如果需要重新计算梯度，记得清零 x.grad.zero_() y.backward() print(x.grad) # tensor([2.0, 4.0, 6.0]) ``` **⚠️ 重要提醒**：PyTorch官方文档指出，**在绝大多数情况下不需要设置 `retain_graph=True`**，而且通常可以用更高效的方式来替代。过度使用 `retain_graph=True` 会导致**内存泄漏**，因为计算图永远不会被释放。 ### 5.3 `create_graph=True`：构建导数计算图用于高阶求导 **`create_graph=True`** 的作用更加深刻：它告诉PyTorch，**在计算梯度的过程中，也要构建计算图**。 这意味着什么？**梯度本身变成了一个可以继续求导的“函数”** ——允许我们计算**高阶导数**（如二阶导数、海森矩阵）。 ```python x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) y = x.pow(2).sum() # 第一次求导：create_graph=True 让梯度本身也带有计算图 grad = torch.autograd.grad(y, x, create_graph=True)[0] print(grad) # tensor([2.0, 4.0, 6.0], grad_fn=<...>) print(grad.requires_grad) # True ← 梯度本身还可以求导！ # 对梯度再求导（二阶导数） z = grad.pow(2).sum() z.backward() print(x.grad) # tensor([8.0, 32.0, 72.0]) # 数学验证：d²/dx² (x²) = 2，所以 d/dx (2x)² = 8x # 对x=[1,2,3]：8*1=8, 8*2=16? 等等... ``` 等等，让我们仔细推导一下： - $y = x_1^2 + x_2^2 + x_3^2$ - $\partial y/\partial x_i = 2x_i$ - 对梯度再求导：$\partial^2 y/\partial x_i^2 = 2$ - 但我们计算的是 $z = \sum_i (\partial y/\partial x_i)^2 = \sum_i 4x_i^2$ - 所以 $\partial z/\partial x_i = 8x_i$，对 $x=[1,2,3]$ 得到 $[8, 16, 24]$ ```python # 修正验证 x.grad.zero_() grad = torch.autograd.grad(y, x, create_graph=True)[0] z = grad.pow(2).sum() z.backward() print(x.grad) # tensor([ 8., 16., 24.]) ✓ ``` ### 5.4 两者的关系与区别总结 | 参数 | 作用 | 默认值 | 典型场景 | |------|------|--------|----------| | `retain_graph` | 反向传播后**保留**计算图 | `False`（释放） | 多次反向传播 | | `create_graph` | 在梯度计算时**构建**计算图 | `False` | 高阶导数（海森矩阵） | **关键区别**： - **`retain_graph`** 控制的是**原计算图**的存留——用于**多次使用同一个图** - **`create_graph`** 控制的是**导数计算图**的构建——用于**对梯度再次求导** ```python # 组合使用：计算高阶导数并多次反向传播 x = torch.tensor([1.0], requires_grad=True) # 第一次：计算一阶导数，同时构建导数图 grad1 = torch.autograd.grad(x**3, x, create_graph=True)[0] # 3x² # 第二次：对一阶导数再求导（二阶导数） grad2 = torch.autograd.grad(grad1, x, create_graph=True)[0] # 6x # 第三次：对二阶导数再求导（三阶导数） grad3 = torch.autograd.grad(grad2, x)[0] # 6 print(grad1, grad2, grad3) # 3.0, 6.0, 6.0 # x³ 的三阶导数 = 6 ✓ ``` ## 六、静态图 vs 动态图 ### 6.1 两种计算图哲学的底层差异 目前主流的深度学习框架在计算图的构建时机上分为两大阵营： - **静态图（Static Graph）** ：**先定义，再运行**（Define-and-Run） - **动态图（Dynamic Graph）** ：**边运行，边定义**（Define-by-Run） **静态图的代表**：TensorFlow 1.x **动态图的代表**：PyTorch ### 6.2 静态图（TensorFlow 1.x）：先画好“施工图纸” 在静态图模式下，**计算图在代码执行之前就被完整地定义好了**。 **工作流程**： 1. **定义阶段**：用框架的API“描述”整个计算图（占位符、操作、损失函数、优化器） 2. **编译阶段**：框架对计算图进行优化（算子融合、内存复用等） 3. **执行阶段**：将数据“喂”给编译好的图，反复执行 ```python # TensorFlow 1.x 风格（静态图） import tensorflow as tf # 阶段1：定义图 x = tf.placeholder(tf.float32, shape=[None, 784]) W = tf.Variable(tf.random_normal([784, 10])) b = tf.Variable(tf.zeros([10])) y = tf.nn.softmax(tf.matmul(x, W) + b) # 阶段2：编译（创建Session） sess = tf.Session() sess.run(tf.global_variables_initializer()) # 阶段3：执行 output = sess.run(y, feed_dict={x: data}) ``` **优点**： - **性能优化**：框架可以在执行前对整个计算图进行全局优化 - **分布式训练**：图可以方便地分布到多台机器上 - **部署方便**：可以将图导出为独立文件（如TensorFlow的`.pb`文件） **缺点**： - **调试困难**：无法在中间插入`print`查看变量值 - **灵活性差**：改变网络结构（如动态分支）非常困难 - **开发体验差**：定义和执行分离，不符合Python的“交互式”习惯 ### 6.3 动态图（PyTorch）：边走边画“地图” 在动态图模式下，**计算图在代码运行时实时构建**。**每一次前向传播都会构建一个新的计算图**。 ```python # PyTorch 风格（动态图） import torch x = torch.randn(32, 784) W = torch.randn(784, 10, requires_grad=True) b = torch.randn(10, requires_grad=True) # 图在运行时构建！ y = torch.softmax(x @ W + b, dim=1) loss = y.sum() loss.backward() # 反向传播，然后图被释放 ``` **优点**： - **易于调试**：可以在任何位置插入`print`查看张量值 - **灵活性强**：支持动态控制流（if、for循环可以根据数据改变结构） - **符合直觉**：代码执行顺序就是图的构建顺序 **缺点**： - **性能开销**：每次都要重新构建图，无法进行全局优化 - **部署稍复杂**：需要额外的序列化机制（如TorchScript） ### 6.4 现代框架的融合趋势 **静态图和动态图之争的本质是“开发效率”与“运行性能”的权衡**。 近年来，两大阵营正在**相互融合**： - **PyTorch 2.0** 通过 `torch.compile()` 实现**动态图到静态图的自动转换**，性能接近TensorFlow静态图 - **TensorFlow 2.x** 默认启用**动态图模式（Eager Execution）** ，同时保留 `@tf.function` 进行静态图优化 | 特性 | 静态图（TF1.x） | 动态图（PyTorch） | 现代趋势 | |------|----------------|-------------------|----------| | 图构建时机 | 执行前 | 运行时 | 混合模式 | | 调试体验 | 差 | 好 | 动态为主 | | 运行性能 | 优 | 中 | 动态+编译优化 | | 灵活性 | 差 | 优 | 动态为主 | | 部署 | 方便 | 需转换 | 各自优化 | > 有趣的是，PyTorch的动态图机制在实现上有一个形象的称呼——**Wengert List**（或tape）。它本质上是**一个按顺序记录所有操作的动态数据结构**，反向传播时从后往前“回放”这个列表来计算梯度。 ## 七、实战：手写一个微型自动微分框架 为了真正理解自动微分的底层原理，我们来**从零实现一个极简的自动微分系统**。 ### 7.1 设计思路 我们的微型框架需要支持： 1. **张量（Tensor）** ：存储数值和梯度 2. **计算图节点**：记录操作历史（`grad_fn`） 3. **反向传播**：沿着计算图逆向传播梯度 ```python import math class Tensor: """极简张量类，支持自动微分""" def __init__(self, data, children=(), op=''): self.data = data if isinstance(data, (int, float)) else float(data) self.grad = 0.0 self._backward = lambda: None # 默认无操作 self._prev = set(children) # 前驱节点 self._op = op # 操作名称 def __add__(self, other): other = other if isinstance(other, Tensor) else Tensor(other) out = Tensor(self.data + other.data, (self, other), '+') def _backward(): self.grad += out.grad other.grad += out.grad out._backward = _backward return out def __mul__(self, other): other = other if isinstance(other, Tensor) else Tensor(other) out = Tensor(self.data * other.data, (self, other), '*') def _backward(): self.grad += other.data * out.grad other.grad += self.data * out.grad out._backward = _backward return out def __pow__(self, other): assert isinstance(other, (int, float)), "只支持数值指数" out = Tensor(self.data ** other, (self,), f'**{other}') def _backward(): self.grad += other * (self.data ** (other - 1)) * out.grad out._backward = _backward return out def backward(self): """拓扑排序后反向传播""" # 构建拓扑排序 topo = [] visited = set() def build_topo(v): if v not in visited: visited.add(v) for child in v._prev: build_topo(child) topo.append(v) build_topo(self) # 反向传播 self.grad = 1.0 for v in reversed(topo): v._backward() ``` ### 7.2 测试我们的微型框架 ```python # 测试：计算 y = (x + 2) * 3，其中 x = 1 x = Tensor(1.0) a = x + 2 y = a * 3 y.backward() print(f"dy/dx = {x.grad}") # 3.0 ✓ # 测试：计算 y = x²，其中 x = 3 x = Tensor(3.0) y = x ** 2 y.backward() print(f"dy/dx = {x.grad}") # 6.0 ✓ # 测试：多变量 y = x₁² + x₂² x1 = Tensor(2.0) x2 = Tensor(3.0) y = (x1 ** 2) + (x2 ** 2) y.backward() print(f"∂y/∂x₁ = {x1.grad}") # 4.0 ✓ print(f"∂y/∂x₂ = {x2.grad}") # 6.0 ✓ ``` ### 7.3 这个微型框架揭示了什么？ 通过这个简单的实现，我们可以清晰地看到自动微分的本质： 1. **每个操作都记录了自己的前驱节点**（`_prev`）和**反向传播函数**（`_backward`） 2. **前向传播**时，每个节点只计算自己的值，并构建计算图 3. **反向传播**时，通过**拓扑排序**确定计算顺序，从输出向输入逐层调用 `_backward` 4. 每个节点的 `_backward` 函数封装了**该操作对应的链式法则** 这其实就是PyTorch Autograd的**微型版本**——区别只在于PyTorch的实现更加高效、支持GPU、处理了更多边界情况。 ## 八、总结 回顾本文的核心脉络： 1. **计算图**是自动微分的“骨架”——用有向无环图表示计算过程，前向传播按拓扑顺序计算，反向传播按逆序传播梯度。 2. **链式法则**是反向传播的“数学引擎”——在计算图中，**一条路径的偏微分等于路径上各相邻节点偏微分的连乘**，多路径则求和。 3. **雅可比矩阵**描述向量函数的**一阶**敏感性（$m \times n$），**海森矩阵**描述标量函数的**二阶**偏导（$n \times n$）。PyTorch的Autograd默认计算的是**向量-雅可比积（VJP）** ，而非完整矩阵。 4. **`retain_graph=True`** 保留原计算图供多次反向传播；**`create_graph=True`** 为梯度构建计算图以支持**高阶导数**。两者在默认情况下均为 `False`。 5. **静态图**（TF1.x）先定义后执行，性能优但调试难；**动态图**（PyTorch）边运行边构建，灵活易调试但有一定性能开销。现代框架正在**融合两者**——PyTorch 2.0的 `torch.compile()` 和TensorFlow 2.x的Eager Execution都是明证。 6. 通过**手写微型自动微分框架**，我们揭示了Autograd的底层本质——**每个操作都记录前驱节点和反向传播函数，反向传播时通过拓扑排序逐层调用**。 理解自动微分与计算图的底层原理，不仅能帮助你更好地使用PyTorch/TensorFlow等框架，更能让你在遇到**梯度异常、显存不足、高阶导数需求**等实际问题时，具备**从底层诊断和解决问题的能力**。这，正是一个优秀AI工程师的核心竞争力所在。 Autograd & Computational Graph —— 自动微分与计算图
/posts/ai_math/autograd_computational_graph/ 分享：
Some information may be outdated