跳到主要内容

数据截至 (上游 commit 7bc720e951fe)

01 · engine.py 自动微分引擎

这一章讲什么: 整个 micrograd 的心脏——micrograd/engine.py 这 94 行。读完你会能徒手讲清:计算图什么时候长出来、梯度凭什么能「自动」流回去、拓扑排序在这里到底解决什么问题。


1. 它要解决的小问题

给定一个标量表达式,比如 g = (a*b + b**3 - a) ** 2 / 2,我们想要 dg/da、dg/db。

手工做法是「数值微分」:把 a 扰动一点看 g 变多少。但它有两个死穴:

  • :有 n 个输入就至少算 2n 次前向。神经网络有百万参数,不可行。
  • 不准:浮点下扰动步长很难选,误差不可控。

正确做法是反向模式自动微分(reverse-mode autodiff):利用链式法则,从输出往回走一趟,一次算出所有输入的精确梯度。深度学习里的「反向传播」就是它。

问题是:这趟「往回走」在代码里怎么组织?需要存什么、按什么顺序走?——这就是 engine.py 回答的事。


2. 思路:节点存三样东西,反向就是调闭包

把表达式拆成一步一步的初等运算(a*bb**3、…),每步产生一个节点。关键洞察:链式法则的每个因子都是「局部的」——z = x + y 时 dz/dx = 1、dz/dy = 1;z = x * y 时 dz/dx = y、dz/dy = x。每条局部法则只依赖这一步的两个操作数,与表达式其余部分无关。

所以一个节点只需存三样东西:

存的字段是什么Value.__init__
data本节点的前向数值micrograd/engine.py:6
grad最终输出对本节点的导数(初始 0,反传时填)micrograd/engine.py:7
_backward本节点的局部链式法则:「我收到自己的 grad 后,往各输入的 grad 上加多少」micrograd/engine.py:9
_prev / _op由哪些节点、经哪个算子算出(构图与画图用)micrograd/engine.py:10-11

前向每算一步,顺手把这一步的 _backward 定义好。反向时:把输出节点的 grad 设为 1(∂g/∂g=1),然后保证每个节点都在它的消费者之后被处理,逐个调 _backward——每个节点的 grad 在被使用的那一刻必然是完整的,局部法则一乘,梯度就流回了叶子。

「保证处理顺序」靠的就是拓扑排序,见 §4。


3. 图示:一张小图的前向与反向

e = a * b + c 为例(a=2, b=3, c=4)。左到右是前向,反传时按编号 ④③②① 的顺序倒着走:

a=2 ─┐
├─► ① d = a*b = 6 ──┐
b=3 ─┘ ├─► ③ e = d+c = 10 ──► (输出)
c=4 ──────────────────────┘

反向(grad 初始全 0,e.grad 置 1):
③ e._backward: d.grad += 1·1, c.grad += 1·1 (加法:梯度原样分流)
① d._backward: a.grad += b.data·d.grad = 3
b.grad += a.data·d.grad = 2 (乘法:梯度乘对方数值)
收工:a.grad=3(∂e/∂a=b),b.grad=2,c.grad=1 —— 与手算一致

怎么读这张图: 前向沿箭头建节点;反向沿箭头倒走,每个节点把自己那份 grad 按局部法则「分账」给来源。一个变量若同时喂给多个下游(比如 c + c),它会收到多笔分账——这正是 += 存在的理由,见 §6。


4. 原理演示:20 行写完整个引擎

下面这段把 engine.py 的两个核心机制(闭包存法则 + 拓扑排序)压到最小,帮助建立直觉:

# 示意,非源码
class V:
def __init__(self, data, children=()):
self.data, self.grad = data, 0.0
self._backward = lambda: None # 默认:叶子不需要反传
self._prev = set(children)

def __mul__(self, o):
out = V(self.data * o.data, (self, o))
def _backward(): # 闭包捕获 self/o/out 三个局部量
self.grad += o.data * out.grad # dz/dx = y,乘以来路梯度
o.grad += self.data * out.grad
out._backward = _backward
return out

def backward(root):
topo, seen = [], set()
def build(v): # 深度优先,先子后父
if v not in seen:
seen.add(v)
for c in v._prev: build(c)
topo.append(v)
build(root)
root.grad = 1.0 # ∂root/∂root = 1
for v in reversed(topo): # 倒序 = 消费者先于生产者
v._backward()

重点看两件事:

  • _backward前向那一刻就定好的——它捕获的 o.dataself.data 就是当时的数值。建图与算值同时发生,这是动态图。
  • reversed(topo) 是反传正确性的全部保证:一个节点被调 _backward 时,它的 grad 一定已被所有下游累加完。

5. 真实实现

5.1 算子与它的局部法则:__mul__

micrograd/engine.py:24-33,符号 Value.__mul__:

def __mul__(self, other):
other = other if isinstance(other, Value) else Value(other)
out = Value(self.data * other.data, (self, other), '*')

def _backward():
self.grad += other.data * out.grad
other.grad += self.data * out.grad
out._backward = _backward

第一行把裸数字包装成 Value,所以 2 * aa * 2a * b 三种写法都成立(配合 __rmul__,micrograd/engine.py:84-85)。_backward 两行就是乘法的链式法则:dz/dx = y,dz/dy = x。

5.2 加法与 ReLU

Value.__add__(micrograd/engine.py:13-22)的 _backwardself.grad += out.grad——加法是「梯度分流器」,原样拷给两个输入。

Value.relu(micrograd/engine.py:45-52)的 _backwardself.grad += (out.data > 0) * out.grad——ReLU 在正区间导数为 1、负区间为 0,用 out.data > 0 这个布尔值当开关。

5.3 反向驱动器:backward()

micrograd/engine.py:54-70,符号 Value.backward:

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
for v in reversed(topo):
v._backward()

build_topo 是标准的深度优先拓扑排序:递归处理完所有 _prev 孩子才 append 自己,所以 topo生产者在前、消费者在后;reversed 之后就是反传需要的「消费者先清账」顺序。self.grad = 1 是边界条件:输出对自己的导数是 1。

5.4 只写四个算子,其余全靠归约

真正带 _backward 的算子只有四个:__add____mul____pow__relu。其余运算全部是组合(micrograd/engine.py:72-91):

便利运算归约成位置
-xx * -1micrograd/engine.py:72-73
x - yx + (-y)micrograd/engine.py:78-79
x / yx * y**-1micrograd/engine.py:87-88

这是代码量能压到 94 行的关键手法:需要维护局部导数法则的表面积被缩到最小,新增便利运算永远不碰反向逻辑。


6. 关键细节与坑

  1. 为什么是 += 而不是 = 若一个节点被多个下游使用(如 y = x + x,或 demo 里 c += c + 1),它的总梯度是所有下游贡献之和(多元链式法则的加法)。= 会只保留最后一笔,+= 才对。所有四个算子的 _backward 都遵守这一点(如 micrograd/engine.py:18-19)。PyTorch 语义相同。
  2. += 的代价:必须手动清零。 因为梯度是累加的,第二次 backward() 会叠在上一次的值上。所以训练循环里每一步都要先 zero_grad()(micrograd/nn.py:6-8)——这不是样板代码,是累加语义的必然配套。
  3. __pow__ 只支持常数指数。 micrograd/engine.py:36assert isinstance(other, (int, float)),所以 a ** b(b 是 Value)会直接报错;x**2 可以。教学取舍:省掉「指数也是变量」的通用幂法则。
  4. _prev 是 set,顺序不影响正确性。 拓扑序只要求「孩子在父前」,同一层的先后无所谓;set 去重顺带处理了「同一节点被引用多次」。
  5. 没有「反向图」这种东西。 PyTorch 文档里常说的 backward graph,在这里并不存在为独立数据结构——图只有一张(前向建的那张),反向只是按序触发闭包。理解这一点,就祛魅了自动微分的一大半。
  6. grad 也存在 Value 上,但不参与建图。 _backward 里写的是裸浮点运算(other.data * out.grad),不会再生成新节点——也就是说这套引擎天然是「一阶」的,不支持梯度再求导(inferred:闭包内只读 .data,未包装成 Value 运算)。