Autograd 不是魔法,它只是链式法则的工程化。 数学上一行公式,实现上需要回答三个问题: 谁是谁的输入?局部梯度怎么算?多路径梯度怎么合并? 理清这三件事,backward 算法就只剩下"按图反着走一遍"。

三个核心原语的局部梯度

链式法则是数学,要把它跑起来还得有承载图结构的数据结构。 我们设计一个 Value 类,它既存储数据, 又是计算图的一个节点——前向时被链入图,反向时被遍历到。 只需要 4 个字段 + 1 个闭包。

Value 类骨架

class Value:
    """标量自动微分节点。data: float, grad: float, _children: tuple[Value]."""
    def __init__(self, data, _children=(), _op=''):
        self.data       = data           # 前向值
        self.grad       = 0.0            # 累积梯度,初始 0
        self._children  = _children      # 反向遍历用的"边"
        self._op        = _op            # 仅用于调试 / 打印
        self._backward  = lambda: None  # 默认 no-op,叶子节点用

    def __repr__(self):
        return f"Value(data={self.data:.4f}, grad={self.grad:.4f}, op={self._op!r})"

数据结构准备好,backward 算法只有两个步骤: ① 用 DFS 把整个 DAG 拓扑排序;② 按逆序遍历,每个节点调一次 _backward()。 为什么拓扑排序?因为必须先收齐所有上游贡献的梯度,再往下游分发—— 否则一个节点的 .grad 还不完整就用它做局部反向,结果就错了。

backward 的 12 行实现

def backward(self):
    """从 self 出发触发反向传播,梯度回填到所有可达的叶节点。"""
    # 1) DFS 后置顺序构造拓扑序(从叶到根)
    topo, visited = [], set()
    def build(v):
        if v in visited: return
        visited.add(v)
        for child in v._children:
            build(child)
        topo.append(v)
    build(self)

    # 2) 自己的梯度初始化为 1 (默认 dL/dL = 1)
    self.grad = 1.0

    # 3) 反向遍历,逐节点调用其闭包
    for v in reversed(topo):
        v._backward()

把前三节合起来,完整代码不到 200 行。 支持 + - * / ** exp log relu 等基本 op, 足以训练一个手写的 MLP—— 这就是 Karpathy micrograd 的核心思想。

Value 类全实现

# micrograd.py — 标量自动微分,纯 Python ~150 行
from __future__ import annotations
import math

class Value:
    def __init__(self, data, _children=(), _op=''):
        self.data       = float(data)
        self.grad       = 0.0
        self._children  = _children
        self._op        = _op
        self._backward  = lambda: None

    # ----- 二元运算 -----
    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')
        def _backward():
            self.grad  += 1.0 * out.grad
            other.grad += 1.0 * out.grad
        out._backward = _backward
        return out

    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   # ∂out/∂self  = other
            other.grad += self.data  * out.grad   # ∂out/∂other = self
        out._backward = _backward
        return out

    def __pow__(self, k):                  # 仅支持常数指数
        assert isinstance(k, (int, float))
        out = Value(self.data ** k, (self,), f'**{k}')
        def _backward():
            self.grad += (k * self.data ** (k - 1)) * out.grad
        out._backward = _backward
        return out

    # ----- 一元非线性 -----
    def relu(self):
        out = Value(max(0.0, self.data), (self,), 'relu')
        def _backward():
            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad
        out._backward = _backward
        return out

    def exp(self):
        out = Value(math.exp(self.data), (self,), 'exp')
        def _backward():
            self.grad += out.data * out.grad   # d(eˣ)/dx = eˣ
        out._backward = _backward
        return out

    def log(self):
        out = Value(math.log(self.data), (self,), 'log')
        def _backward():
            self.grad += (1.0 / self.data) * out.grad
        out._backward = _backward
        return out

    # ----- 派生运算 (复用上面已有的) -----
    def __neg__(self):           return self * -1
    def __sub__(self, o):        return self + (-o)
    def __rsub__(self, o):       return o + (-self)
    def __truediv__(self, o):    return self * (o ** -1 if isinstance(o, (int, float)) else o ** -1)
    def __radd__(self, o):       return self + o
    def __rmul__(self, o):       return self * o

    # ----- backward -----
    def backward(self):
        topo, visited = [], set()
        def build(v):
            if id(v) in visited: return
            visited.add(id(v))
            for c in v._children:
                build(c)
            topo.append(v)
        build(self)
        self.grad = 1.0
        for v in reversed(topo):
            v._backward()

    def __repr__(self):
        return f"Value(data={self.data:.4f}, grad={self.grad:.4f}, op={self._op!r})"

最小冒烟测试

from micrograd import Value

a = Value(2.0)
b = Value(-3.0)
c = Value(10.0)
e = a * b
d = e + c          # d = a·b + c = -6 + 10 = 4
f = d.relu()       # f = relu(4) = 4
f.backward()

print(a.grad, b.grad, c.grad)
# 期望:
#   ∂f/∂a = b * 1 (relu 通过) = -3
#   ∂f/∂b = a * 1 = 2
#   ∂f/∂c = 1
# 输出:-3.0  2.0  1.0  ✓

用它训练一个 MLP

import random
from micrograd import Value

# --- 用 Value 搭建一个最小神经网络 ---
class Neuron:
    def __init__(self, nin):
        self.w = [Value(random.uniform(-1, 1)) for _ in range(nin)]
        self.b = Value(0.0)
    def __call__(self, x):
        act = sum((wi * xi for wi, xi in zip(self.w, x)), self.b)
        return act.relu()
    def parameters(self):
        return self.w + [self.b]

class Layer:
    def __init__(self, nin, nout):
        self.neurons = [Neuron(nin) for _ in range(nout)]
    def __call__(self, x):
        return [n(x) for n in self.neurons]
    def parameters(self):
        return [p for n in self.neurons for p in n.parameters()]

class MLP:
    def __init__(self, sizes):
        self.layers = [Layer(sizes[i], sizes[i+1]) for i in range(len(sizes)-1)]
    def __call__(self, x):
        for l in self.layers: x = l(x)
        return x[0] if len(x)==1 else x
    def parameters(self):
        return [p for l in self.layers for p in l.parameters()]

# --- 数据:学一个简单 XOR-like 数据集 ---
xs = [[2.0,3.0], [3.0,-1.0], [0.5,1.0], [1.0,1.0]]
ys = [1.0, -1.0, -1.0, 1.0]

model = MLP([2, 4, 4, 1])

# --- 训练循环 ---
for step in range(200):
    # 1) 前向 + loss
    preds = [model(x) for x in xs]
    loss  = sum((p - y)**2 for p, y in zip(preds, ys))

    # 2) 清梯度 + 反向
    for p in model.parameters(): p.grad = 0.0
    loss.backward()

    # 3) SGD 更新
    lr = 0.05
    for p in model.parameters():
        p.data -= lr * p.grad

    if step % 20 == 0:
        print(f"step {step:3d}  loss = {loss.data:.4f}")

# 你应该看到 loss 一路下降到接近 0,200 行代码就训练了一个 MLP

手写的代码要和参考实现交叉验证才算可信。 把同一个表达式分别在 micrograd 和 PyTorch 里算一遍, 对比梯度——一致就说明你的实现是对的。 这也是写新 op 时最有效的回归测试方法。

import torch
from micrograd import Value

# ----- 测试函数:f(a, b, c) = (a * b + c).relu() ** 2 -----

# 1) micrograd 版本
a1 = Value(2.0); b1 = Value(-3.0); c1 = Value(10.0)
y1 = ((a1 * b1 + c1).relu()) ** 2
y1.backward()
mg = (y1.data, a1.grad, b1.grad, c1.grad)

# 2) PyTorch 版本
a2 = torch.tensor(2.0,  requires_grad=True)
b2 = torch.tensor(-3.0, requires_grad=True)
c2 = torch.tensor(10.0, requires_grad=True)
y2 = ((a2 * b2 + c2).relu()) ** 2
y2.backward()
pt = (y2.item(), a2.grad.item(), b2.grad.item(), c2.grad.item())

# 3) 对比
for name, x, y in zip(["y", "a.grad", "b.grad", "c.grad"], mg, pt):
    print(f"{name:8s}  micrograd={x:8.4f}  pytorch={y:8.4f}  ok={abs(x-y)<1e-6}")

# 期望输出:
#   y         micrograd= 16.0000  pytorch= 16.0000  ok=True
#   a.grad    micrograd=-24.0000  pytorch=-24.0000  ok=True
#   b.grad    micrograd= 16.0000  pytorch= 16.0000  ok=True
#   c.grad    micrograd=  8.0000  pytorch=  8.0000  ok=True

两边的对应关系

推荐扩展练习

1. 手推 z = (a + b * c) ** 2 对 a / b / c 的偏导,写出每个中间变量的局部梯度,标出反向遍历顺序。
2. 写一个 tanh 算子加入 micrograd:实现 _backward,并用 PyTorch 对照验证(梯度公式 1 - tanh²(x))。
3. 解释为什么反向模式 AD 的总耗时与前向同阶,但前向模式 AD 需要 N 次前向才能算所有 N 个输入的梯度。
4. 把 micrograd 的递归 DFS 改成迭代版本(用显式栈)。运行你的 MLP demo 验证结果一致。
5. 设计实验:把 y = x * x 的 backward 跑一遍,然后故意把 += 改成 =,观察 x.grad 错在什么数字上,解释为什么。