为什么需要自动微分

深度学习训练的核心循环可以概括为一句话：前向传播算损失，反向传播算梯度，优化器更新参数。其中「反向传播算梯度」这一步，正是由自动微分引擎（Autograd）完成的。

考虑一个典型的神经网络损失函数：

$$L(\theta) = \sum_i \ell\bigl(f(x_i; \theta),\, y_i\bigr)$$

其中$\theta $可能包含数亿个参数。训练时我们需要求出$\frac{\partial L}{\partial \theta }$ 的每一个分量——对 175B 参数的 GPT-3 而言，意味着每次迭代要计算 1750 亿个偏导数。如果靠人来手工推导梯度公式并编码实现，既容易出错又不可维护。自动微分引擎解决的正是这个问题：你只需用普通代码写出前向计算（loss 怎么算出来），Autograd 自动、精确、高效地算出所有梯度。

这份「自动化」背后没有任何魔法，只有两个数学工具：基本运算的求导规则（加、乘、sin、exp……每个都已知）和链式法则。

三种求导方法对比

在自动微分之前，人类已经发明了两种让计算机求导的方法。理解它们的局限，才能理解 Autograd 为什么是「终极答案」。

数值微分：简单但不精确

利用导数定义做近似计算，中心差分公式：

$$ f'(x) \approx \frac{ [f(x+h) – f(x-h)]}{2h}$$

它有两大硬伤：一是精度困境——h 太大截断误差（$O(h^2)$）大，h 太小浮点舍入误差大，实际精度通常只有$10^{-7}~10^{-8}$ 量级，且每个参数需要额外 2 次前向计算，对亿级参数完全不可行。它的唯一价值是作为「标准答案」用来验证 Autograd 的正确性（gradient check）。

符号微分：精确但会爆炸

像数学家一样对表达式做符号变换，例如 SymPy 对 $ \sin(x^2)\cdot \exp(x) $ 求导得到解析表达式。问题在于表达式膨胀（expression swell）：重复子表达式在求导过程中被不断复制，表达式规模指数级增长；更重要的是，它无法自然处理程序中的循环、条件分支、内存操作——而真实神经网络全是这些结构。

自动微分：精确且高效

自动微分的思路完全不同：不对公式求导，对程序的执行过程求导。任何程序无论多复杂，执行时都是一条基本运算的序列（加、乘、矩阵乘、激活函数……），而每个基本运算的导数规则是已知的，最后用链式法则把局部导数串起来。

维度 数值微分 符号微分 自动微分
精度 近似（受浮点限制） 精确 精确（机器精度）
计算复杂度 每个参数 2 次前向 取决于表达式大小 ≈ 1 次前向的常数倍
支持循环/分支 支持 困难 天然支持
代表工具 finite diff SymPy PyTorch / JAX / TF

计算图：Autograd 的骨架

自动微分的实现载体是计算图（Computational Graph）：把程序执行过程表达为有向无环图（DAG）

  • 节点（Node）：输入变量或一个运算的结果（中间变量）
  • 边（Edge）：运算之间的数据依赖关系

下面这个例子贯穿全文：L = (x·y) · max(x·y, z)，取 x=2, y=3, z=4。注意 x·y 被使用了两次，在图中体现为两条出边——这正是计算图表达共享子结构的方式，反向传播时会自动处理「多路径梯度累加」。

关键观察：x 同时影响 L 的两条路径（直接相乘、经过 max），∂L/∂x 是两条路径梯度之和（21 = 18 + 3）。这就是「梯度累加」规则，也是计算图比手写求导更可靠的原因。

前向传播与反向传播

自动微分按梯度传播方向分为两种模式，它们的核心区别在于链式法则的展开顺序。

前向模式（Forward Mode）

沿计算方向同步传播导数。计算节点值的同时，把 $\dot{x}  = \frac{\partial x}{\partial w} $（对某个固定输入 w 的导数）一起算出来。一次前向遍历得到一个输入对所有输出的导数。适合输入少、输出多的场景（参数量小、函数输出多），在控制论、灵敏度分析中常用。

反向模式（Reverse Mode）

先完整做一次前向计算并记录计算图，然后从输出（loss）出发，逆着图的方向，用链式法则逐节点回传梯度$\bar{x}  = \frac{\partial L}{\partial x}$。一次反向遍历得到所有输入对一个输出的导数。

为什么深度学习选反向模式？神经网络的结构恰好是「参数多（百万~千亿）、输出少（一个标量 loss）」。反向模式对这种 多对一 结构一次回传就拿到全部参数梯度，计算成本仅为前向的 2~3 倍。这就是「反向传播算法」的数学本质：反向模式的自动微分。

反向模式的代价：记忆瓶颈

反向模式需要保存前向计算的所有中间变量（因为反向计算局部导数时要用到前向值，如 $\frac{\partial L}{\partial x} = b$ 需要知道 b 的值）。这带来了深度学习最著名的工程问题——显存占用。激活检查点（activation checkpointing，用时间换显存）、梯度检查点等技术正是为了缓解这个问题（见第七节）。

逐步推漾示例

以图 2 为例，反向传播从 $\frac{\partial L}{\partial x} = 1$出发，按拓扑逆序逐节点计算：

注：a 有两条入边路径，此处按梯度累加规则合并计算，最终 ∂L/∂x = 21、∂L/∂y = 14，与直接展开解析式一致。max 的局部导数是「赢家通吃」：梯度 100% 流向较大的输入，另一个输入梯度为 0——这就是 max、ReLU 等函数的反向规则。

PyTorch Autograd 架构剖析

有了前面的理论基础，我们来看工业级实现。PyTorch 的 Autograd 是一个精心设计的 C++ 系统，自上而下分为四层：

Tensor 与 AutogradMeta：梯度的家

在 PyTorch 0.4 之后，Variable 被合并进 Tensor。每个 Tensor 内部挂载一个可选的 AutogradMeta 结构，核心字段：

  • grad_：梯度存放地，调用.backward() 后在此累积
  • grad_fn_：指向创建该 Tensor 的反向节点（AddBackward、MulBackward……）
  • requires_grad_：是否追踪该 Tensor 的运算
import torch

x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([3.0], requires_grad=True)
z = torch.tensor([4.0], requires_grad=True)

a = x * y          # grad_fn=MulBackward0
b = torch.max(a, z) # grad_fn=MaxBackward0
L = a * b          # grad_fn=MulBackward0

L.backward()       # 启动反向传播

print(x.grad, y.grad, z.grad)
# tensor([21.]) tensor([14.]) tensor([0.])
# 与图 2 的手工推演完全一致

动态图：边执行边建图

PyTorch 是动态图（define-by-run）框架：前向传播执行到哪里，计算图就即时构建到哪里。每个运算经过 Dispatcher 时，AutogradKernel 会：

  • 检查输入是否有requires_grad=True，没有则直接跳过建图（快速路径）
  • 调用真正的后端 kernel 完成数值计算
  • 创建对应的反向Node，通过 Edge 连接到输入的 grad_fn，挂到输出的 AutogradMeta 上

这让计算图天然支持 Python 的 if/for/while——每次迭代的图都可以不同，这也是 PyTorch 在研究界流行的原因之一。

Autograd Engine：反向图的并行执行器

调用 loss.backward() 时，引擎从 loss 节点出发，把所有反向 Node 按拓扑逆序调度执行。工程亮点：

  • 每个设备一个线程：同一设备的反向节点串行（保证写安全），不同设备（多 GPU）并行回传
  • ready_queue + 引用计数：一个 Node 的梯度依赖全部到齐（依赖计数归零）才入队执行，天然实现拓扑排序与多路径梯度累加
  • CUDA 流同步：反向 kernel 提交到与前向相同的流，保证执行顺序

自定义算子：torch.autograd.Function

当需要引入框架没有的操作（或 fused kernel 优化）时，手动声明前向与反向：

class MyReLU(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)   # 保存反向所需的中间量
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors     # 取回前向值
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0    # 局部导数：赢家通吃
        return grad_input

注意：backward 里写入的是 grad_output（上游回传的梯度），返回的是对各输入的 grad_input。引擎负责沿 Edge 把这些梯度继续传给上游节点并自动累加——你只需要关心「局部导数」这一步。

动手实现一个迷你 Autograd

理解 Autograd 最好的方式是亲手写一个。下面约 60 行 Python 实现了完整的反向模式自动微分（Value 类方案，致敬 micrograd）：

class Value:
    """标量自动微分：data=前向值, grad=反向梯度, _backward=局部导数闭包"""

    def __init__(self, data, parents=()):
        self.data = data
        self.grad = 0
        self._backward = lambda: None
        self._parents = parents      # (父节点, 局部梯度) 列表

    def __add__(self, other):
        out = Value(self.data + other.data, ((self, 1.0), (other, 1.0)))
        def _backward():
            self.grad += 1.0 * out.grad   # 加法：梯度直通
            other.grad += 1.0 * out.grad
        out._backward = _backward
        return out

    def __mul__(self, other):
        out = Value(self.data * other.data,
                    ((self, other.data), (other, self.data)))
        def _backward():
            self.grad += other.data * out.grad   # 乘法：交换乘
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

    def relu(self):
        out = Value(max(0, self.data), ((self, 0.0),))
        def _backward():
            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad
        out._backward = _backward
        return out

    def backward(self):
        # 1. 拓扑排序（DFS）
        topo, visited = [], set()
        def build(v):
            if v not in visited:
                visited.add(v)
                for p, _ in v._parents:
                    build(p)
                topo.append(v)
        build(self)
        # 2. 从输出反向传播，梯度累加
        self.grad = 1.0
        for v in reversed(topo):
            v._backward()


x, y, z = Value(2), Value(3), Value(4)
a = x * y
b = Value(max(a.data, z.data), ((a, 1.0 if a.data > z.data else 0.0),
                                (z, 1.0 if z.data >= a.data else 0.0)))
L = a * b
L.backward()
print(x.grad, y.grad, z.grad)   # 21 14 0 —— 与图 2 一致

这个迷你实现包含了工业级 Autograd 的全部核心思想：

迷你实现 PyTorch 对应 作用
_parents 列表 Edge + Node 记录反向图拓扑
_backward 闭包 grad_fn.apply() 封装局部导数计算
grad += 引擎的 Buffer 累加 多路径梯度累加
DFS 拓扑排序 ready_queue 依赖计数 确定反向执行顺序

高级话题

动态图 vs 静态图

维度 动态图（PyTorch） 静态图（TF1 / JAX jit）
建图时机 执行时即时建图 先定义完整图，再执行
Python 控制流 原生支持，每步可变 需 trace / tf.cond 等机制
调试体验 可直接 print / pdb 需特殊的调试工具
编译优化 需 torch.compile 补足 天然利于算子融合、XLA

PyTorch 2.x 的 torch.compile 通过 TorchDynamo 捕获动态图并做算子融合，正在融合两者的优点。

高阶梯度：create_graph=True

把反向传播本身也建到计算图里，就能对梯度再求梯度（二阶导），用于 penalty 正则（如 WGAN-GP）、meta-learning（MAML）等场景：

g = torch.autograd.grad(loss, x, create_graph=True)[0]
gg = torch.autograd.grad(g.sum(), x)[0]   # 二阶导

显存优化：no_grad 与 checkpoint

  • no_grad()：推理/验证时关闭建图，省显存且更快
  • utils.checkpoint：前向不保存中间激活，反向时重新计算（时间换显存，可训练更深的模型）
  • backward()` 后紧跟 `optimizer.zero_grad()：PyTorch 梯度默认累加（支持梯度累积训练），不用时须手动清零

常见陷阱

现象 原因 解法
loss 不下降、grad 为 None 叶子节点未设 requires_grad 或图被 no_grad 截断 检查建图路径
显存持续增长 保留了带图的历史输出（如每步 loss.append(loss)） 存 loss.item()
RuntimeError: backward twice 默认图在 backward 后释放 设 retain_graph=True
原地操作报错 in-place 修改破坏了反向所需的值 改用 out-of-place 写法

总结

回顾全文的主线：

  • 自动微分的本质：把程序分解为基本运算序列，用已知的局部导数 + 链式法则自动合成完整梯度——精确到机器精度，成本仅为前向的常数倍。
  • 计算图是骨架：节点存值与反向闭包，边表达依赖；多路径梯度累加由图结构自动处理。
  • 反向模式是深度学习的必然选择：多参数 → 单 loss 的结构下，一次反向遍历拿到全部梯度。
  • PyTorch 实现：Tensor 挂载 AutogradMeta，动态执行时经 Dispatcher 记录反向节点，Autograd Engine 以拓扑序多线程执行反向闭包。
  • 60 行代码即可复现其核心机制——Autograd 不是魔法，是链式法则的工程化。

一句话总结：Autograd = 计算图（记录）+ 链式法则（数学）+ 拓扑排序执行引擎（工程）。深刻理解这三件事，就理解了整个深度学习框架的心脏。

0
相关文章:
  1. 开源因果推断库CausalML
  2. 基于网格的经纬度轨迹聚类
  3. 纯 Python 从零实现 GPT
  4. JavaScript学习之版本特性
  5. Python学习之面向对象基础