从零实现自动微分框架

反向传播(Backpropagation)是一种用于训练神经网络的算法。它依赖自动微分技术,自动微分是一种计算导数的技术。他们是深度学习的底层关键。

读完上面这段话,你看懂了吗?坦白地讲,我是一点都没看懂

幸运的是,Andrej Karpathy 大神出马了。他开发了 micrograd 教学框架,并录制了高质量的教学视频,全面讲解这背后是怎么回事。在大神的号召下,互联网涌现出一大片延伸的教学资料。

随着学的深入我,我发现只从代码上理解,并不能满足我,我更希望从数学理论上得到透彻理解。因此我将著名的《花书》也纳入在内,文章的前半部分重点梳理理论,后半部分介绍 micrograd。

在这么多资料的加持下,再加上还有一众 AI 导师,只要我学,我就一定能会!因此在本文中,我将从一头雾水开始,逐步吃透自动微分技术。

网络资料

本文建立在大量教学资料的肩膀上,我梳理了一部分,请直接滑动到文末进行浏览。

由于本文是边学边完善,阅读起来略显破碎,可直接参见网络资料,他们比我写的更好。比如,《从0实现一个极简的自动微分框架 - Nagi-ovo》一文,已经完成了本文所要完成的目标。

但本文对于我来说是重要的。因为网上的资料都是别人的知识。在完成本文后,知识将真正装进我的脑子里。

同《从零训练GPT》一文一样,我会尽可能详尽的查找认知的盲区,如同穿珍珠一般,将知识串联起来。


深度学习的训练目标

在深度学习中,我们经常需要训练具有大量参数的复杂模型,如 GPT 等大语言模型。这些模型本质上是一个多层复合函数 F(x;θ),其中 x 表示输入(如文本),而 θ 表示模型参数(如神经网络的权重和偏置)。训练的目标是找到最优参数 θ,使得模型在给定任务上的损失函数 L 最小化:

θ=argminθL(F(x;θ),y)
注意

模型本身是一个非线性函数,但是我们并不是直接对模型本身求极值。而是我们会引入一个损失函数/目标函数。模型是一个概率密度函数,我们传入输入数据后,得到一个先验估计,同时,我们有对应输入的实际输出。损失函数是对先验估计和实际输出,通过某种运算,得到的误差而衡量。而损失函数,此时是模型参数的函数,我们对它求取极值,即确定能让误差最小的模型参数。

其中 y 表示期望输出(如文本标签)。为了使用梯度下降法等优化算法求解最优参数,我们需要计算损失函数 L 相对于每个参数 θi 的梯度 Lθi。这就是自动微分的任务。

假设我们正在训练一个语言模型,如 GPT,其目标是在给定前几个单词的情况下预测下一个单词。我们可以将这个模型表示为一个函数 F(x;θ),其中 x 是输入的单词序列,而 θ 是模型的参数(如transformer块中的权重和偏置)。

我们的目标是找到最优参数θ,使得模型在训练数据上的损失函数 L 最小化:

θ=argminθL(F(x;θ),y)

其中 y 是真实的下一个单词。

在这个优化问题中,损失函数 L 定义了一个超曲面。

我们的目标是找到这个超曲面上的全局最小值点 θ,即损失函数最小的参数值。然而,这个超曲面通常非常复杂,可能存在多个局部最小值、鞍点和平坦区域,使得优化变得困难。

这时,反向传播和自动微分就派上用场了。它们就像一个训练有素的向导,引导我们沿着超曲面向下走,朝着全局最小值前进。具体来说:

  1. 正向传播:我们首先从起点出发,将训练数据 x 输入到模型中,计算损失函数 L。这相当于在超曲面上选择一个起点。

  2. 反向传播:然后,我们利用反向传播算法,自动计算损失函数 L 在该点,相对于每个参数 θi 的梯度 Lθi。这个梯度告诉我们,在当前点,沿着每个参数方向走一小步,损失函数会如何变化。它就像一个指南针,指引我们下一步该往哪个方向走。

  3. 梯度下降:最后,我们根据梯度的方向,更新每个参数:

    θiθiαLθi

    其中 α 是学习率,控制每次更新的步长。这相当于沿着梯度的反方向,在超曲面上走一小步,从而使损失函数减小。

我们反复执行这个过程,每次都让参数沿着梯度下降的方向前进一点,直到达到一个最小值点(希望是全局最小值)。这就像一个下山的过程。

在这个过程中,反向传播负责计算梯度,告诉我们下一步该往哪里走。而自动微分则负责高效实现反向传播,自动计算复杂模型的梯度,而无需手动推导和编码。它们合作无间,成为深度学习的核心动力

它们自动高效地计算梯度,引导模型参数沿着最陡下降方向前进,最终达到损失函数的最小值点。这就像一个智能向导,带领我们在复杂的参数空间中寻找最优解。没有它们,训练大型神经网络将变得举步维艰。


Jacobian 雅可比矩阵

通常,模型的输入和输出都是向量,我们需要计算所有的偏导数。包含所有这样的偏导数的矩阵,称为雅可比矩阵(Jacobian matrix),用来表示其梯度。

考虑一个函数 F:RnR,它将一个 n 维向量 x=(x1,x2,,xn) 映射为一个标量输出 y。在神经网络的上下文中,x 通常表示网络的参数(权重和偏置),而 y 表示网络的输出,例如图像中狗出现的概率。

神经网络通常由许多子层组成,因此我们可以将函数 F 看作是多个原始函数(如加法、乘法、激活函数等)的复合:

F(x)=f4f3f2f1(x)

其中f1,f2,f3,f4表示神经网络中的四个原始函数。

现在,让我们来计算函数 F 相对于其输入变量 x 的梯度。这个梯度实际上是一个雅可比矩阵,它由函数 F 相对于每个输入变量的偏导数组成:

JF(x)=yx=[yx1yx2yxn]

由于函数 F 只有一个标量输出,因此其雅可比矩阵是一个1×n的行向量。如果输出也是向量,则雅可比矩阵是一个矩阵

在实现自动微分时,我们需要高效地计算这个雅可比矩阵。这可以通过反向传播算法来实现,该算法利用链式法则,以递归的方式计算复合函数的梯度。


停一停,捋一捋

到目前为止,我明白了:

  1. 神经网络的函数表示

    • 神经网络是一个非线性函数 F(x;θ),在推理阶段,输入数据 x 是自变量,参数 θ 是常量。
    • 在训练阶段,输入 x 和输出 y 是常量,参数 θ 是自变量。
  2. 损失函数

    • 损失函数 L 是关于参数 θ 的函数,我们需要最小化损失函数来找到最佳参数。
  3. 导数和雅可比矩阵

    • 对于单变量函数,求导得到导数。当输入是向量时,对每个分量求导,得到的矩阵形式叫做雅可比矩阵,用来表示梯度。
  4. 数值微分

    • 雅可比矩阵只是理论上表示了梯度,那具体如何计算呢?
    • 可以使用数值微分的方法。具体来说,给定一组初始参数 θ,第一步是给第一个参数一个小的增量,观察损失函数的变化,从而得到损失函数对该参数的偏导数。对其他参数依次进行相同操作,得到所有参数的偏导数。
    • 但这种方法效率很低,因为每次参数更新都需要重新计算整个模型的推理过程,对于大规模模型,这是不可行的。
    • 注意,自动微分与数值微分是有区别的:自动微分通过解析计算导数,精度高且效率高;数值微分通过有限差分近似计算导数,容易受舍入误差影响且计算效率低。
  5. 反向传播算法

    • 为了提高效率,使用反向传播算法。反向传播利用链式法则,将模型的求导过程分解成多个步骤。在计算图中,许多参数间的计算存在重叠,因此可以复用中间过程的梯度。这大大减少了计算量,从而提高了运算效率。
  6. 自动微分

    • 反向传播只是将梯度计算分解成多个小步骤,但最终仍然需要计算偏导数。这时就需要用到自动微分技术。自动微分是一种在程序中自动计算函数导数的技术,通过将计算图中的每个操作记录下来,然后反向计算每个操作的导数,从而得到最终的梯度。
    • 自动微分是一种计算导数的方法,与数值微分(通过微小增量计算导数)和符号微分(通过解析公式计算导数)不同,自动微分通过链式法则递归地计算复合函数的导数。它在处理复杂函数时具有高效和精确的优点。

前向传播算法(Forawrd Propagation)

以一个简单的三层MLP为例:

graph LR
    x[输入层] --> h1[隐藏层1] 
    h1 --> h2[隐藏层2]
    h2 --> y[输出层]
    y --> L[损失函数]

前向传播计算每层的输出:

h1=f1(x;W1,b1)
h2=f2(h1;W2,b2)
y=f3(h2;W3,b3)

前向传播值得就是推理,传入一个输入 x,经过每一层运算,最终在输出层得到 y^,然后与真实值 y 一起传入损失函数,得到最终的损失值。

这个过程就是前向传播。


计算图(Computational Graph)

前向传播和反向传播,看似是对称关系,其实也不是。前向传播是求值的,反向传播是根据新的损失函数的信息,反向流动,来更新梯度

在继续探讨反向传播之前,我们需要先理解一个重要的概念:计算图。

计算图是一种用于表示计算过程的数据结构,它由节点和有向边组成。在计算图中:

下面是一个简单的计算图示例:

graph LR
    x[x] --> mul(("×"))
    w[w] --> mul
    mul --> add(("+"))
    b[b] --> add
    add --> y[y]

在这个例子中,输入变量 x 和权重 w 相乘,然后加上偏置 b,得到输出 y。可以用数学公式表示为:y=wx+b

计算图清晰地描述了数据和操作之间的依赖关系。这种表示方式有几个重要的优点:

  1. 自动微分:通过跟踪计算图中的操作,可以自动计算出任意变量的梯度,而无需手动推导梯度公式。这是反向传播算法的基础。

  2. 高效计算:计算图揭示了操作之间的并行性,可以充分利用现代硬件(如GPU)进行加速。

  3. 可视化:计算图提供了一种直观的方式来理解复杂的计算过程,特别是对于深度神经网络。

实际上,PyTorchTensorFlow 等深度学习框架的核心,就是一个自动构建和执行计算图的引擎。理解计算图的概念,对于深入理解这些框架的工作原理至关重要。

在下一节中,我们将基于计算图,详细探讨反向传播算法的原理和实现。通过沿着计算图反向传播梯度,我们就可以高效地训练神经网络模型。


反向传播算法 (Back Propagation)

反向传播是一种用于高效计算神经网络参数梯度的算法。它基于链式法则,通过递归地计算从输出到每个参数的梯度,使得我们能够使用梯度下降等优化算法来训练神经网络

计算图与链式法则

让我们以一个简单的计算图为例来理解反向传播的原理。考虑如下的计算图:

graph LR
    x((x)) --> a((a))
    a --> b((b))
    b --> c((c))
    c --> d((d))
    d --> e((e))
    e --> L((L))

其中,x 是输入,L 是我们关心的输出(如损失函数)。每个节点表示一个操作,例如:

a=f(x)=2x
b=g(a)=a2
c=h(b)=3b
d=i(c)=sin(c)
e=j(d)=ed
L=k(e)=log(e)

我们的目标是计算 L 相对于 x 的梯度 Lx。根据链式法则:

Lx=Leeddccbbaax

反向传播算法

反向传播算法就是以一种巧妙的方式组织链式法则中的局部梯度计算和梯度传递。具体来说,它分为两个阶段:

  1. 前向传播:按照拓扑序逐个计算每个节点的输出值。
  2. 反向传播:从输出节点 L 开始,按照拓扑序的逆序,递归地计算每个节点的梯度,并将其传递给前驱节点。

在反向传播阶段,每个节点 v 需要完成两个任务:

  1. 计算 v 的输出相对于其输入的局部梯度
  2. v 的梯度传递给其前驱节点

例如,对于节点e,我们首先计算 Le=1e(这里LL初始化为1)。然后,将这个梯度传递给节点d

Ld=Leed=1eed=ed1

依此类推,直到我们计算出Lx


micrograd

micrograd 是一个用于自动微分神经网络训练的小型库,用于教学目的。由 Andrej Karpathy 开发。

该库最大的特点是小——小于 100 行 Python 代码。

micrograd 本质上是一个 AutoGrad(Automatic Gradient)引擎。它的功能是实现反向传播

反向传播这个算法能让我们有效评估一个关于权重的神经网络,在损失函数下,让我们可以迭代式地调优,调整神经网络的权重,使得损失函数最小化。这样就可以提高网络的准确性。反向传播是现代神经网络库(如 PyTorch)的数学核心。

micrograd 是一个标量的自动微分引擎,标量就是一个数值。作为教学目的,易于理解。在实际中,处理的是多维的张量,张量就是这些标量的数组,利用计算机的并行计算能力,高效计算。

核心代码只有两个:engine.py 自动微分引擎,对神经网络一无所知。nn.py:建立在 micrograd 引擎上的神经网络库。

micrograd 是训练神经网络所必要的核心原理,其它的所有附加都是为了提升运算效率。


micrograd example

在 micrograd 的首页中,给出了一段示例代码:

from micrograd.engine import Value

a = Value(-4.0)
b = Value(2.0)
c = a + b
d = a * b + b**3
c += c + 1
c += 1 + c + (-a)
d += d * 2 + (b + a).relu()
d += 3 * d + (b - a).relu()
e = c - d
f = e**2
g = f / 2.0
g += 10.0 / f

# prints 24.7041, the outcome of this forward pass
print(f'{g.data:.4f}') 
g.backward()

# prints 138.8338, i.e. the numerical value of dg/da
print(f'{a.grad:.4f}') 

# prints 645.5773, i.e. the numerical value of dg/db
print(f'{b.grad:.4f}') 

在这段代码中,创建了一个表达式,有两个输入 a 和 b,经过一系列运算,围绕 a 和 b 创建了一个表达式图,最终输出为 g。

a 和 b 以==值对象(Value)==的形式创建,分别是 -4 和 2。值对象(Value)是 micrograd 的关键,在目前,它的作用是存储对应的数字值。在后续小节中,将对其进行扩充。

在上述表达式中包含一些数学运算,比如加法和乘法,它们是 micrograd 支持的运算,在后面小节中将以运算符重载的方式,分别实现。

对于这个简单运算,为什么要用值对象封装?原因是,通过使用值对象,micrograd 会在背后建立整个数学表达式,它会存储每个值的值、运算类型、到上级运算数(如 a 和 b 之于 c)的指针,在未来还会存储梯度。

使用正推法,即带入 a 和 b 的值执行运算,得到结果 g 的值。用 g.data 即可获得正推法结果的值。

g.backward() 会初始化在节点 g 处的反向传播。反向传播会从 g 开始,往这个数学表达式图的反方向,然后递归应用微积分链式法则。这就能让我们计算,g 关于整个表达式图中对所有节点的导数。所有节点比如 b、d、c 等,还有 a 和 b。

a.gradb.grad 分别查询了 g 关于 a 和 b 的导数。这个导数非常重要,它会告诉我们 a 和 b 对 g 的影响。这两个导数告诉我们,如果 a 和 b 往正向增大一点点,g 会如何变化

神经网络只是一系列数学表达式,就像上面一样。它们将输入的数据作为输入,取某个神经网络的权重作为输入,经过数学表达式,输出神经网络的预测,或者说损失函数。


导数的数值计算方法

对于如下函数,如何计算对 f 对于 x 的导数?

def f(x):
	return 3*x**2 - 4*x + 5

按照微分的定义,取很小的 delta(h),计算函数斜率:

h = 0.00001
x = 3.0
(f(x + h) - f(x)) / h

注:h 不能取得过小,超出浮点数能够处理的边界,这会导致运算错误。

该方法仅用于演示如何以数值方法求微分。在 micrograd 的实际实现中,并不是这么求微分的。在 micrograd 实现中,通过链式法则,将微分拆解到最小粒度,如加减乘除,由于加减乘除的微分是可以推导出来的,因此直接用导出公式计算即可。

我想,对于那些不易求出导数解析公式的函数,此处的数值计算方法就有用武之地了。


偏导数的计算方法

上一节中介绍了单变量情况,在本节中考虑多变量。如下代码中,有 3 个标量 a、b、c。如何计算 d 对于 a 的偏导数呢?

a = 2.0
b = -3.0
c = 10.0
d = a * b + c
print(d) # 4.0

计算对 a 的偏导数:

h = 0.0001

a = 2.0
b = -3.0
c = 10.0

d1 = a * b + c
a += h
d2 = a * b + c

print('d1', d1)
print('d2', d2)
print('slope', (d2 - d1) / h)

跟上一节一样,给 a 添加一点点增量,计算出斜率。(注:计算对 b、c 的偏导数也是同理。)


Micrograd Value 值对象

Micrograd 的核心对象是 Value 对象。它是一个具有值和梯度的简单对象。Micrograd 的 engine.py 里只有一个类,就是 Value。这里直接给出 Vaule 的完整实现(省略部分非核心代码,比如运算符):

class Value:
    """ stores a single scalar value and its gradient """

    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0
        # internal variables used for autograd graph construction
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op # the op that produced this node, for graphviz / debugging / etc

    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 += out.grad
            other.grad += out.grad
        out._backward = _backward

        return out

	# ...

    def backward(self):

        # topological order all of the children in the graph
        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)

        # go one variable at a time and apply the chain rule to get its gradient
        self.grad = 1
        for v in reversed(topo):
            v._backward()

	# ...

    def __repr__(self):
        return f"Value(data={self.data}, grad={self.grad})"

在上述代码中,首先是之前小节介绍过的,使用 data 表示数值。


加法运算符

在上述代码中,Value 实现了一系列运算符(上述代码中,我只保留了加法)。以加法 __add__ 为例:other 表示加数,将相加的和用一个新 Value 实例存放。同时,将被加数和加数,作为 Value 的 _children,这样便建立起指针引用关系。

Value 类的 __add__ 方法中,我们定义了一个 _backward 函数,用于计算加法操作的局部梯度。

首先,回顾一下 __add__ 方法的实现:

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 += out.grad
        other.grad += out.grad
    out._backward = _backward

    return out

当我们执行 a + b 操作时,会创建一个新的 Value 对象 out,它的 data 属性是 a.data + b.data,同时将 ab 作为 out 的前驱节点。

现在,让我们考虑反向传播过程。假设我们已经计算出了 out 的梯度 out.grad,那么根据链式法则,我们需要将这个梯度传递给 ab

对于加法操作,我们知道:

out=a+bouta=1,outb=1

这意味着,out 的梯度对 ab 的梯度贡献是相同的,都等于 out.grad。如何理解呢?

当我们说 outa=1,我们确实意味着 out 的任何微小变化都会引起 a 相同的变化量,因为 aout 的影响是直接且完全的。同样,outb=1 表明 bout 的影响也是直接且完全的。

在反向传播中,已经计算出的 out.grad 表示 out 的微分变化量。因为加法的性质是每个加数直接全量影响结果,所以:

这表明 ab 都将完整地、独立地接收到从 out 传递过来的梯度信息,每个变量的梯度都是 out 的梯度。

直觉理解:

画一个简单的图表:设 x 轴代表 a,y 轴代表 b,z 轴代表 out。现在,想象一个通过 x 轴和 y 轴的原点斜率为 1 的平面。这个平面表示了加法 out = a+b 的关系。如果 z 增加了一定的量 out.grad(斜率),则无论是沿 x 轴的 a 增加还是沿 y 轴的 b 增加,都需要相应地增加 out.grad。

想象你和一个朋友各自推着一个等重的物体上坡。如果坡度突然增加(这类似于 out.grad 的增加),为了维持相同的上升速度,你和你的朋友都需要增加相同的力量。在这里,坡度的增加是 out.grad,而你和你的朋友增加的力量就是 a.grad 和 b.grad。


理解 _backward 方法

在上一节中,我们在 Value 类中引入了 _backward 方法,用于计算局部梯度并将梯度传递给前驱节点。现在,让我们深入理解这个方法的作用和实现。

__add__ 方法的实现中,_backward 函数的作用是计算加法操作的局部梯度,并将其传递给操作数 selfother。根据加法的梯度规则,我们知道:

outself=1,outother=1

这意味着,out 的梯度对 selfother 的梯度贡献是相同的,都等于 out.grad。因此,在 _backward 函数中,我们将 out.grad 累加到 self.gradother.grad 中。

现在,让我们考虑反向传播的过程。当我们调用 out.backward() 时,会发生以下步骤:

  1. out.grad 初始化为 1。
  2. 调用 out._backward(),计算加法操作的局部梯度,并将其传递给 selfother
  3. 递归地调用 self.backward()other.backward(),继续反向传播过程。

通过这种方式,梯度信息就能够沿着计算图的反方向流动,直到到达输入节点。

需要注意的是,_backward 方法只定义了局部梯度的计算规则,而实际的梯度传播过程是在 backward 方法中完成的。backward 方法会按照拓扑序的逆序遍历计算图,并依次调用每个节点的 _backward 方法


乘法运算符

掌握加法之后,接下来看乘法运算符。乘法的梯度规则与加法有所不同,需要使用乘法法则:

ifz=xy,thenzx=y,zy=x

这意味着,如果输出 out 是两个输入 ab 的乘积,那么在反向传播过程中:

为什么说,这里是乘积的关系?

当我们处理神经网络或其他包含多步运算的计算模型时,通常会涉及多个变量和多个步骤的复合函数。在这种情况下,链式法则允许我们正确地计算一个变量的变化如何影响整个表达式的输出。

链式法则的基本思想是,如果你有一个由多个函数组成的复合函数,你可以通过将这些函数的导数(梯度)相乘来找到最终输出对任何一个输入的敏感度。这就像是在说,要知道最终输出如何对一个内部步骤的改变做出反应,你需要考虑这个步骤到最终输出之间所有步骤的"影响力"。

在你的例子中,设想 z=ab,并且 z 是某个更大函数的一部分,这个大函数的输出我们叫做 out。如果我们想知道 aout 的影响,我们首先需要知道 zout 的影响,这就是所谓的 out 的梯度。

这里的关键是理解 zout 的贡献是如何通过 ab 的乘积传递的。根据乘法规则,za 的敏感度是 b。但是,我们还需要知道 z 的这一变化如何影响最终的 out

所以,如果 zout 的梯度是某个值(我们称之为 outz),并且我们知道 za 的梯度是 b,那么 aout 的总梯度将是这两者的乘积:

outa=outz×za

因此,我们有:

outa=outz×b

这意味着,a 的梯度是 b 乘以 zout 的梯度。这样的乘法展示了 a 的一个小变动是如何通过 z 这一中间变量影响 out 的。

我悟了!

原来,对每个输入节点(ab),算得是 outa,对此应用链式法则。这样彻底解答了为何这里是乘积的关系。

代码实现如下:

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

	return out

其中:

  1. 首先,我们重载 __mul__ 运算符,使其能够处理 Value 对象之间的乘法。

  2. 我们创建一个新的 Value 对象 out,其值为两个输入的乘积,同时将操作数 selfother 设置为 out 的前驱节点,并将操作类型设置为 '*'

  3. 然后,我们定义 _backward 方法,用于计算乘法操作的局部梯度。根据乘法的梯度规则:

    • self 的梯度为 other.data * out.grad
    • other 的梯度为 self.data * out.grad
  4. 最后,我们将 _backward 方法赋给 out._backward,以便在反向传播过程中调用。


拓扑逆序排序

在前面的小节中,我们已经知道,反向传播需要按照特定的顺序遍历计算图,以正确地计算和传播梯度。这个特定的顺序被称为拓扑逆序Topological Reverse Order)。

所谓拓扑逆序,是指按照节点之间的依赖关系,从输出节点开始,逆向遍历整个计算图的顺序。直觉上理解,就是从计算图的末端开始,逐步向前,直到到达输入节点。

为什么要按照拓扑逆序遍历计算图呢?这是因为,根据链式法则,一个节点的梯度计算依赖于其后继节点的梯度。换句话说,只有当一个节点的所有后继节点的梯度都计算完毕,我们才能计算这个节点自身的梯度。

举个例子,考虑如下的计算图:

graph LR
    a((a)) --> c((c))
    b((b)) --> c
    c --> d((d))

在这个图中,节点 d 赖于其前驱节点 c,而节点 c 又依赖于其前驱节点 ab。因此,正确的反向传播顺序应该是:d -> c -> ab。这就是拓扑逆序。

在 micrograd 的实现中,我们可以使用深度优先搜索DFS)算法来实现拓扑排序。具体来说,就是从输出节点开始,递归地遍历其前驱节点,并将节点按照完成遍历的逆序排列。下面是一个简化的实现:

def topological_sort(node):
    visited = set()
    topo_order = []

    def dfs(v):
        if v not in visited:
            visited.add(v)
            for child in v._prev:
                dfs(child)
            topo_order.append(v)

    dfs(node)
    # 将 `topo_order` 列表中的元素顺序颠倒过来。
    # 第一个冒号表示从头到尾进行切片(不指定起始和结束位置)。
    # 第二个冒号和 `-1` 表示步长为 `-1`,即从列表的末尾开始,
    #                 步长为 `-1` 表示逆向遍历整个列表。
    return topo_order[:: -1]

在这个实现中,我们使用一个 visited 集合来记录已经遍历过的节点,避免重复遍历。我们从输出节点开始,递归地调用 dfs 函数遍历其前驱节点。当一个节点的所有前驱节点都遍历完毕时,我们将其添加到 topo_order 列表的末尾。最后,我们将 topo_order 逆序,得到拓扑逆序的节点列表。

有了拓扑逆序,我们就可以按照正确的顺序进行反向传播了。在下一小节中,我们将看到如何利用拓扑逆序,在 backward 函数中实现自动梯度计算。


实现反向传播

在本节中,我们将详细探讨如何在 micrograd 中实现反向传播算法。反向传播是自动微分的核心,它利用链式法则,高效地计算复杂计算图中每个节点的梯度

首先,让我们回顾一下前面的内容。我们已经了解了如何使用计算图来表示复杂的数学表达式,以及如何通过前向传播来计算表达式的值。现在,我们需要找到一种方法,来自动计算每个节点(即每个 Value 对象)的梯度。

反向传播算法分为两个主要步骤:

  1. 对计算图进行拓扑排序。
  2. 按照拓扑序的逆序,递归地计算每个节点的梯度。

有了拓扑序,我们就可以按照正确的顺序计算每个节点的梯度了。这个过程是递归进行的,从输出节点开始,沿着计算图的反方向,应用链式法则,直到到达输入节点。

下面是 backward 方法的实现:

def backward(self):

	# topological order all of the children in the graph
	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)

	# go one variable at a time and apply the chain rule to get its gradient
	self.grad = 1
	for v in reversed(topo):
		v._backward()

这个方法首先将输出节点的梯度设置初始值为 1(这是反向传播的起点)。然后,它按照拓扑序的逆序遍历计算图,根据每个节点的操作类型(+*),使用链式法则计算其前驱节点的梯度,并将其累加到前驱节点的 grad 属性中。

这就是 micrograd 中反向传播算法的核心实现。它自动、高效地计算了复杂表达式中每个变量的梯度,为优化神经网络的参数提供了基础。


处理计算图中的节点复用

在某些情况下,计算图中的节点可能会被多条路径复用。这意味着,同一个节点可能会出现在计算图的不同分支中,从而对最终输出产生多次影响。在反向传播过程中,我们需要正确地处理这种节点复用的情况,以确保梯度计算的正确性。

让我们考虑下面这个计算图的例子:

graph LR
    a((a)) --> b((b))
    a --> c((c))
    b --> d((d))
    c --> d

在这个图中,节点 a 同时影响了节点 b 和节点 c,而节点 d 则依赖于节点 b 和节点 c。换句话说,节点 a 通过两条不同的路径影响了节点 d

那么,在反向传播过程中,我们该如何处理这种情况呢?关键在于正确地累加来自不同路径的梯度贡献。

具体来说,当我们计算节点 a 的梯度时,我们需要考虑它对节点 d 的两个贡献:一个通过节点 b,另一个通过节点 c。根据链式法则,节点 a 的梯度应该是这两个贡献的总和:

da=dbba+dcca

在 micrograd 的实现中,这个梯度累加的过程是自动完成的。回顾一下我们的 __add____mul__ 方法:

def __add__(self, other):
    # ...
    def _backward():
        self.grad += out.grad
        other.grad += out.grad
    # ...

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

注意,在这两个方法的 _backward 函数中,我们都是使用 += 运算符来更新 self.gradother.grad。这意味着,如果一个节点通过多条路径接收到梯度,这些梯度会被自动累加起来。

所以,在反向传播过程中,当我们按照拓扑逆序遍历计算图时,复用的节点会自然而然地接收到来自不同路径的梯度贡献,并将它们累加起来,得到正确的总梯度。

这种自动梯度累加的机制,使得我们可以构建复杂的计算图,而无需手动跟踪每个节点的多个梯度贡献。这大大简化了自动微分的实现,并提高了其适用性和可扩展性。


进军神经网络:模块基类

在深入探索 micrograd 如何构建神经网络之前,让我们先来谈谈软件设计中的一个重要原则:模块化。模块化设计的目标是将一个复杂系统分解为多个独立、可重用的模块,每个模块负责一个特定的功能。这种分而治之的方法有几个重要的优势:

  1. 提高代码的可读性和可维护性。
  2. 使代码更容易测试和调试。
  3. 提高代码的可重用性,减少重复工作。
  4. 使系统更容易扩展和修改。

micrograd 同样遵循这一设计原则。在 nn.py 文件中,作者定义了一个名为 Module 的基类:

class Module:

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

    def parameters(self):
        return []

这个 Module 类看似简单,但它为构建更复杂的神经网络模块奠定了基础。让我们仔细看看它的两个方法:

  1. zero_grad():这个方法用于将模块中所有参数的梯度重置为0。这通常在每次反向传播之前调用,以清除上一次迭代的梯度。

  2. parameters():这个方法返回一个包含模块所有参数的列表。在基类中,它返回一个空列表,因为基础的 Module 类并没有任何参数。但在子类中,我们会重写这个方法,返回子类特有的参数列表。

有了这个基类,我们就可以通过继承和组合的方式,构建出各种复杂的神经网络模块,如全连接层、卷积层、循环神经网络等。每个模块都有自己的参数和前向传播逻辑,但它们都共享相同的接口(如 zero_grad()parameters()),这使得它们可以无缝地组合在一起。


Neuron:神经网络的基本构建块

在上一节中,我们介绍了 micrograd 中的 Module 基类,它为构建复杂的神经网络模块提供了基础。现在,让我们来看看如何使用这个基类来实现神经网络的基本构建块:神经元neuron)。

nn.py 文件中,作者定义了一个名为 Neuron 的类,它继承自 Module 基类:

class Neuron(Module):

    def __init__(self, nin, nonlin=True):
        self.w = [Value(random.uniform(-1,1)) for _ in range(nin)]
        self.b = Value(0)
        self.nonlin = nonlin

    def __call__(self, x):
        act = sum((wi*xi for wi,xi in zip(self.w, x)), self.b)
        return act.relu() if self.nonlin else act

    def parameters(self):
        return self.w + [self.b]

    def __repr__(self):
        return f"{'ReLU' if self.nonlin else 'Linear'}Neuron({len(self.w)})"

让我们逐步分析这个类的实现:

  1. __init__ 方法:这是神经元的构造函数。它接受两个参数:

    • nin:表示神经元的输入数量。
    • nonlin:一个布尔值,表示是否在神经元的输出上应用 ReLU 激活函数。默认为 True

    在构造函数中,我们初始化了神经元的权重 self.w 和偏置 self.b。权重是一个长度为 nin 的列表,每个权重都是一个 Value 对象,其初始值随机采样自区间 [-1, 1]。偏置也是一个 Value 对象,初始值为0。

  2. __call__ 方法:这个方法定义了神经元的前向传播逻辑。当我们将一个输入向量 x 传递给神经元时,这个方法会被调用。它计算神经元的激活值 act,即输入向量 x 与权重 self.w 的点积,再加上偏置 self.b。如果 self.nonlinTrue,则在激活值上应用 ReLU 函数;否则直接返回激活值。

  3. parameters 方法:这个方法返回神经元的所有参数,即权重列表 self.w 和偏置 self.b。这个方法覆盖了基类的同名方法。

  4. __repr__ 方法:这个方法提供了一个字符串表示,用于打印或显示神经元对象。根据 self.nonlin 的值,它返回 "ReLUNeuron(n)" 或 "LinearNeuron(n)",其中 n 是输入的数量。

这就是 micrograd 中的 Neuron 类。它封装了一个单个神经元的计算逻辑,包括权重、偏置、前向传播和参数访问。通过组合多个神经元,我们就可以构建出更复杂的神经网络。

理解

神经元接收多个输入(长度为 nin),神经元内部有 nin 个权重参数,用于对输入加权,同时还有一个偏置参数。神经元的正向传播,就是对输入进行加权运算。如果是非线性的,再将这个运算结果输入 ReLU 非线性函数。这里面每一步运算都是基于 Value 层层封装的,因此是可以反向传播计算梯度的。


ReLU 运算符

在上一节中,我们探讨了神经元的实现。神经元是神经网络的基本构建块,它接收多个输入,对这些输入进行加权求和,再加上一个偏置项,最后通过一个激活函数产生输出。在 micrograd 中,如果神经元被设置为非线性,那么它将使用 ReLU 作为激活函数。

ReLU,全称为修正线性单元(Rectified Linear Unit),是一种常用的激活函数。它的数学定义如下:

ReLU(x)=max(0,x)

也就是说,ReLU 函数将所有负值输入都置为0,而正值输入则保持不变。这个简单的操作引入了非线性,使得神经网络能够学习和表示复杂的非线性关系。

engine.py 中,ReLU 操作被实现为 Value 类的一个方法:

def relu(self):
    out = Value(0 if self.data < 0 else self.data, (self,), 'ReLU')

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

    return out

这个方法创建了一个新的 Value 对象 out,它的值根据输入 self.data 的正负而定:如果 self.data 小于0,则 out.data 为0;否则,out.data 等于 self.data。同时,这个方法也定义了 ReLU 操作的反向传播规则。

在反向传播过程中,ReLU 函数的梯度计算如下:

ReLU(x)x={0,if x<01,if x0

这意味着,如果一个神经元的输入是正值,那么在反向传播时,它将接收来自后续节点的完整梯度;但如果其输入是负值,那么梯度将被截断为0。这个计算过程在 _backward 函数中通过 (out.data > 0) * out.grad 实现。

ReLU 的这种"选择性"传递梯度的特性,使得基于 ReLU 的神经网络在实践中表现出良好的训练效果和收敛速度。它在一定程度上缓解了梯度消失的问题,因为正值输入的梯度可以无衰减地传播。同时,ReLU 的计算也非常高效,这使其成为最常用的激活函数之一。


Layer:神经网络的组合与堆叠

在前面的章节中,我们已经了解了神经网络的基本构建块:神经元。但是,单个神经元的能力是有限的。为了构建强大的神经网络,我们需要将多个神经元组合起来,形成层(Layer)。

在 micrograd 中,Layer 类同样继承自 Module 基类:

class Layer(Module):

    def __init__(self, nin, nout, **kwargs):
        self.neurons = [Neuron(nin, **kwargs) for _ in range(nout)]

    def __call__(self, x):
        out = [n(x) for n in self.neurons]
        return out[0] if len(out) == 1 else out

    def parameters(self):
        return [p for n in self.neurons for p in n.parameters()]

    def __repr__(self):
        return f"Layer of [{', '.join(str(n) for n in self.neurons)}]"

让我们逐步剖析这个类的实现:

  1. __init__ 方法:这是层的构造函数。它接受两个主要参数:

    • nin:输入的维度,即前一层的神经元数量。
    • nout:输出的维度,即本层的神经元数量。
    • **kwargs:其他关键字参数,用于传递给神经元的初始化参数(如 nonlin)。

    在构造函数中,我们创建了一个包含 nout 个神经元的列表 self.neurons。每个神经元都接收 nin 维的输入,并根据 **kwargs 进行初始化。

  2. __call__ 方法:这个方法定义了层的前向传播逻辑。给定一个输入向量 x,我们将其传递给层中的每个神经元,并收集它们的输出。如果层中只有一个神经元,则直接返回其输出;否则,返回一个包含所有神经元输出的列表

  3. parameters 方法:这个方法返回层中所有神经元的参数。它使用了两个嵌套的列表推导式:首先遍历层中的所有神经元,然后遍历每个神经元的参数,最后将所有参数收集到一个列表中。

  4. __repr__ 方法:这个方法提供了层的字符串表示,方便打印和调试。它将层中所有神经元的字符串表示用逗号连接起来,形成一个描述性的字符串。

通过组合多个层,我们就可以构建出深度神经网络。例如,一个简单的三层全连接网络可能如下所示:

nin = 2
hidden_size = 4
nout = 1

model = [
    Layer(nin, hidden_size, nonlin=True),
    Layer(hidden_size, hidden_size, nonlin=True),
    Layer(hidden_size, nout, nonlin=False)
]

在这个例子中,我们创建了一个包含两个隐藏层和一个输出层的模型。输入的维度为 2,隐藏层的大小为 4,输出的维度为1。前两个隐藏层使用 ReLU 激活函数,而输出层是一个线性层(不使用激活函数)。

当我们将输入数据传递给这个模型时,数据将依次流经每个层,经过层中的神经元进行变换和处理。最终产生输出。这就是神经网络的前向传播过程。

层的概念使我们能够以一种模块化、可组合的方式构建复杂的神经网络。通过堆叠不同类型和大小的层,我们可以设计出适用于各种任务的网络架构,如图像分类、语言模型、强化学习等。

同时,micrograd 中的 Layer 类也继承了 Module 基类的所有功能,如参数管理(parameters)和梯度清零(zero_grad)。这使得在训练神经网络时,我们可以方便地访问和操作每个层的参数,并在反向传播之前清除上一次迭代的梯度。

总之,Layer 类是 micrograd 中构建神经网络的核心组件之一。它将单个神经元组合成有组织的、可堆叠的单元,使我们能够轻松地设计和训练各种类型的神经网络。


MLP:深度神经网络的简洁实现

在前面的章节中,我们已经了解了 micrograd 中神经元和层的实现。有了这些基本构建块,我们就可以轻松地组合出更复杂的神经网络架构。其中,最经典和常用的一种架构就是==[000.wiki/多层感知机|多层感知机]==。

MLP 由多个全连接层堆叠而成,每个隐藏层使用非线性激活函数(如 ReLU),而输出层通常是一个线性层(不使用激活函数)。这种简单而有效的结构使 MLP 成为许多任务的首选模型,如回归、分类等。

在 micrograd 中,MLP 的实现同样继承自 Module 基类:

class MLP(Module):

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

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

    def parameters(self):
        return [p for layer in self.layers for p in layer.parameters()]

    def __repr__(self):
        return f"MLP of [{', '.join(str(layer) for layer in self.layers)}]"

让我们逐行解析这个简洁而优雅的实现:

  1. __init__ 方法接受两个参数:

    • nin:输入的维度。
    • nouts:一个列表,指定了每个隐藏层和输出层的维度。

    在构造函数中,我们首先将输入维度nin添加到nouts的开头,形成一个完整的层大小列表sz。然后,我们使用一个列表推导式来创建 MLP 的层列表self.layers。对于每两个相邻的层大小sz[i]sz[i+1],我们创建一个Layer对象,并根据层的位置决定是否使用非线性激活函数。最后一层(i == len(nouts)-1)是输出层,不使用激活函数。

  2. __call__ 方法定义了 MLP 的前向传播逻辑。我们将输入x依次传递给每个层,并将每个层的输出作为下一层的输入。最终,我们返回最后一层的输出作为整个 MLP 的输出。

  3. parameters 方法返回 MLP 中所有层的所有参数。这里我们使用了两个嵌套的列表推导式,首先遍历所有层,然后遍历每个层的参数。

  4. __repr__ 方法提供了 MLP 的字符串表示,方便打印和调试。

这个 MLP 类展示了 micrograd 的模块化和可组合性。通过简单地指定输入维度和每个层的输出维度,我们就可以创建任意深度和宽度的多层感知机。例如:

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

这行代码创建了一个具有 3 维输入、两个大小为 4 的隐藏层和一个 1 维输出的 MLP。

在实际使用中,我们可以轻松地将这个 MLP 集成到训练流程中。我们可以使用model.parameters()来访问模型的所有参数,使用model(x)来进行前向传播,并使用model.zero_grad()来清除梯度。这使得 MLP 的训练和使用与其他模块完全一致。


梯度清零:反向传播的重要一步

在上一节中,我们探讨了如何使用 MLP 类来构建和训练深度神经网络。现在,让我们将目光聚焦到训练过程的一个关键步骤:梯度清零。

回想一下,在 micrograd 的 Module 基类中,有一个名为 zero_grad 的方法:

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

这个方法遍历模块中的所有参数,并将它们的梯度设置为 0。但是,为什么我们需要这样做呢?

想象一下,我们正在训练一个神经网络,优化某个损失函数。训练的每一步,我们都会执行以下操作:

  1. 进行前向传播,计算损失函数。
  2. 进行反向传播,计算损失函数对每个参数的梯度。
  3. 使用优化算法(如梯度下降)更新参数。

关键点在第二步。在反向传播过程中,每个参数的梯度是通过累加计算得到的。这意味着,如果我们不在每次迭代之前清零梯度,那么当前步的梯度就会与前一步的梯度累加在一起,导致优化方向的偏差

让我们用一个具体的例子来说明。假设我们有一个参数 w,在第一次迭代中,损失函数对 w 的梯度是 1.0。如果我们不清零梯度,那么在第二次迭代中,即使损失函数对 w 的新梯度是 0.5,w.grad 的值也会是 1.5(1.0 + 0.5),而不是我们想要的 0.5。

因此,在每次迭代之前调用 model.zero_grad() 是至关重要的。这样,我们就能确保每一步的梯度都是基于当前步的损失函数计算得到的,而不会受到前一步梯度的影响。

以下是一个典型的训练循环,展示了梯度清零的位置:

for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        # 清零梯度
        model.zero_grad()
        
        # 前向传播
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
        
        # 反向传播
        loss.backward()
        
        # 更新参数
        optimizer.step()

可以看到,model.zero_grad() 是放在每个批次的训练之前,确保每一步的梯度计算都从零开始

总之,梯度清零是训练神经网络的一个小但重要的步骤。它确保了反向传播过程的正确性,使我们能够有效地优化模型。在 micrograd 中,这个功能被抽象到 Module 基类中,使得所有的模块(如 Layer 和 MLP)都能方便地管理自己的梯度。


网络资源

在自学过程中,我搜集到许多高质量资料,本文是站在这些巨人的肩膀上的,感谢他们:


本文作者:Maeiee

本文链接:从零实现自动微分框架

版权声明:如无特别声明,本文即为原创文章,版权归 Maeiee 所有,未经允许不得转载!


喜欢我文章的朋友请随缘打赏,鼓励我创作更多更好的作品!