跳到主要内容

数据截至 (上游 commit 7bc720e951fe)

02 · nn.py 神经网络库

这一章讲什么: micrograd/nn.py 这 60 行——在 Value 之上搭出「神经元 → 层 → 多层感知机」。它几乎不引入新机制,读完你会明白一个神经网络库的最小充分集到底小到什么程度。

1. 它要解决的小问题​

上一章的引擎只会算 a * b + c 这种散式表达式。要训练,我们需要:

  • 把「一组权重 + 一个偏置 + 激活函数」打包成一个能重复调用的东西(神经元);
  • 把神经元叠成网络;
  • 拿到全部可训练参数的一份清单,好让每个训练步对它们统一做 p.data -= lr * p.grad。

换句话说:给引擎包一层 PyTorch 风格的 Module API。难的不是写,而是看清最少需要包几层、每层负责什么。

2. 思路:组合 Value,递归收集参数​

两个直觉撑起整个文件:

  1. 前向不需要任何新代码。 神经元的计算 w·x + b 再 ReLU,全是 Value 已支持的运算。只要参数本身是 Value,调用网络时计算图就自动长出来——nn.py 完全不用碰 _backward 或拓扑排序。
  2. parameters() 是唯一的正经事。 训练循环需要的只是「所有参数的扁平列表」。于是每层只实现两件事:__call__(前向)和 parameters()(把自己的参数交给上层),父层把子层的列表拼起来。

这也是 PyTorch nn.Module 的同款骨架——micrograd 把它削到了不能再削的形态。

3. 图示:四层抽象的包含关系​

MLP(nin=2, nouts=[16,16,1])
 └─ layers: Layer ×3            第 3 层 nonlin=False(线性输出)
     └─ neurons: Neuron ×nout
         ├─ w: Value ×nin       随机初始化 U(-1,1)
         └─ b: Value            初始 0

parameters() 的收集方向相反:
Neuron 交 w+[b] → Layer 拼 nout 份 → MLP 拼 3 层 → 一份扁平清单

怎么读这张图: 上半是「谁含谁」(构造时自上而下),下半是 parameters() 的汇报链(训练时自下而上)。zero_grad() 走的正是这条汇报链的反向:拿到清单后逐个 p.grad = 0。

4. 原理演示:一个神经元的最小骨架​

# 示意,非源码
class Neuron:
    def __init__(self, nin):
        self.w = [Value(random.uniform(-1, 1)) for _ in range(nin)]  # 每个输入一个权重
        self.b = Value(0.0)                                          # 偏置从 0 出发

    def __call__(self, x):                  # n(x) 直接当前向用
        act = sum((wi*xi for wi, xi in zip(self.w, x)), self.b)  # w·x + b
        return act.relu()                   # 非线性

重点看 __call__ 里没有任何「建图」字样——wi*xi 是 Value.__mul__、sum(..., self.b) 是一串 Value.__add__,图在求值的同时就长好了。nn 库对引擎零感知,这就是分层干净的样子。

5. 真实实现​

5.1 Neuron:参数即 Value​

micrograd/nn.py:15-22,符号 Neuron.__init__ / Neuron.__call__:

self.w = [Value(random.uniform(-1,1)) for _ in range(nin)]
self.b = Value(0)
...
act = sum((wi*xi for wi,xi in zip(self.w, x)), self.b)
return act.relu() if self.nonlin else act

权重初始化是均匀分布 U(-1,1),偏置为 0。sum(...) 的第二个参数 self.b 是起始值——一个小技巧,让 b 也经由 __add__ 进入计算图,反传时自然收到梯度。nonlin=False 时退回线性神经元(输出层用,见 5.3)。

Neuron.parameters(micrograd/nn.py:24-25)就是 return self.w + [self.b]。

5.2 Layer:同构神经元的并列​

micrograd/nn.py:32-37,符号 Layer.__init__ / Layer.__call__:

self.neurons = [Neuron(nin, **kwargs) for _ in range(nout)]
...
out = [n(x) for n in self.neurons]
return out[0] if len(out) == 1 else out

Layer 就是「并排放 nout 个神经元,输出收成一个列表」。唯一的小心机在最后一行:只有一个神经元时直接返回单个 Value 而不是长度为 1 的列表——这让 MLP 的输出层可以当标量分数用,损失函数里写 -yi*scorei 不用先解包。

Layer.parameters(micrograd/nn.py:39-40)用双层推导式把所有神经元的参数拍平。

5.3 MLP:层的串联,最后一层必须线性​

micrograd/nn.py:47-54,符号 MLP.__init__ / MLP.__call__:

sz = [nin] + nouts
self.layers = [Layer(sz[i], sz[i+1], nonlin=i!=len(nouts)-1) for i in range(len(nouts))]
...
for layer in self.layers:
    x = layer(x)

nouts 只给各层宽度,sz 把头尾接上推出每层输入宽度。注意 nonlin=i!=len(nouts)-1:除最后一层外都套 ReLU,最后一层线性——这是分类/回归网络的标配(输出分数不该被 ReLU 截断),在这里用一个布尔表达式就讲完了。

demo 里 MLP(2, [16, 16, 1])(demo.ipynb:96)按此构造:2→16(ReLU)→16(ReLU)→1(线性),共 (2·16+16)+(16·16+16)+(16·1+1)=337 个参数(notebook 输出亦为 337,demo.ipynb:90)。

5.4 Module:zero_grad 与 parameters 的协议​

micrograd/nn.py:4-11,符号 Module:

def zero_grad(self):
    for p in self.parameters():
        p.grad = 0

def parameters(self):
    return []

基类只定义协议:parameters() 默认空,zero_grad() 遍历清单清零。它与上一章「梯度一律 +=」的累加语义配套——每个训练步开始前必须清账,否则梯度跨步累加。PyTorch 的 optimizer.zero_grad() 是同一件事的工程版。

6. 关键细节与坑​

  1. 没有 Tensor、没有批维度。 一次 model(x) 只处理一个样本;demo 里批处理是手写 list(map(model, inputs))(demo.ipynb:127)。这让代码极简,也让它无法谈性能。
  2. x 里的输入最好也是 Value。 Neuron.__call__ 里 wi*xi:若 xi 是裸 float,靠的是 Value.__rmul__(micrograd/engine.py:84-85)兜底,仍能进图;demo 统一 map(Value, xrow) 包一层(demo.ipynb:124),更整齐。
  3. 初始化没有讲究。 U(-1,1) + 零偏置在这个规模能训动,但没有 He/Xavier 初始化——层数一深就会出问题(inferred:无任何按宽度缩放的代码)。教学件不解决这个。
  4. Module 没有参数注册机制。 PyTorch 靠 __setattr__ 拦截自动登记子模块/参数;micrograd 靠每个类手写 parameters() 汇报。少了魔法,多了几行样板——对读懂代码是好事。
  5. nonlin 标志一路透传。 MLP → Layer(**kwargs) → Neuron(micrograd/nn.py:32-33、micrograd/nn.py:49),一个布尔穿过三层,没有配置对象。层数少时这是最简方案。