MicroGrad: 动手实现一个简单的自动微分框架

Micrograd 简介 (Introduction)#

Micrograd 是Andrej Karpathy打造的一个学习性项目，用于理解自动微分、反向传播等基本知识。

视频地址：https://www.bilibili.com/video/BV1De4y1p7Z8

代码地址：https://github.com/karpathy/micrograd

反向传播 (Backpropagation) 是机器学习和深度学习中最常用的一种算法，能够计算神经网络中损失函数相对权重参数的梯度，通过最小化损失在反向传播过程中指导模型权重更新。

导数 (Derivative)#

导数的数学定义#

数学上对导数的定义：函数在某一点上的导数是指函数在这个点附近的变化率，即函数在该点上的切线斜率。公式表示为：

$f'(x) = \lim_{h \to 0} \frac{f(x+h) - f(x)}{h}$

在 x 点处，取一个较小的值 h，如 0.0001，利用上述公式近似计算函数 $f(x)$ 在 $x$ 处的导数。

如：

import math
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

# 定义一个函数
def f(x):
    return 3*x**2 - 4*x + 5
# 使用matplotlib画出函数图像
xs =np.arange(-5, 5, 0.25)
ys = f(xs)
plt.plot(xs, ys)

h = 0.00000001
a = -3.0
(f(a + h) - f(a))/h

输出：-22.00000039920269 就得到了函数 f 在 -3 处的导数近似值

实际上可以通过求导法则直接求出 $f$ 的导数为 $6x - 4$ ，在 - 3 处的导数为 -22
但是在实际的现代神经网络库中不会对每一个函数求出这样的表达式，遇到庞大复杂的函数时求解出这样的导数表达式是困难且低效的。

进一步理解导数#

上文中求解了函数对未知数的导数，但实际的神经网络中，我们需要求解函数对某个标量的导数，这时从导数的定义下手更好理解

h = 0.0001
a = 2.0
b = -3.0
c = 10.0

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

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

按照定义，这里对 c 进行求导（c 进行了少量偏移，求解函数 d 在 c 偏移后的变化率），结果为：
slope 0.9999999999976694

这意味着 c 在增加 0.0001 的微小变化下，d 的变化量约为 c 的变化量的 1 倍（正为增加，副为减少）。

也可以通过导数的定义求解 d 的变化量相对于 c 的变化量的关系：
$\frac {d_2-d_1} {h} = 1$

同理也能求函数相对于 a 和 b 的导数近似值。

神经网络的原理#

以简单的多层感知机 MLP 为例，神经网络实际上是数学公式的集合。数据作为输入，经过隐藏层在神经元中通过数学运算（通常是输入数据的加权求和加上偏置，然后应用非线性激活函数）产生进入到下一层的激活值，最终输出预测结果等其他数据。其主要的运算是一系列的加乘（在线性代数中体现为矩阵运算）以及非线性激活。

前向传播负责计算网络的输出并得到损失值；前向传播完成后执行反向传播，其核心任务是根据损失值和前向传播过程中存储的中间结果，通过链式法则计算损失函数关于各层参数（如权重和偏置）的梯度。通过这些计算得到的梯度，指导隐藏层中神经元的权重参数进行更新，以减小训练误差（其实不仅仅是需要减小训练误差，更重要的是减少测试误差，让模型的泛化性更好才是最终目的）。

在 micrograd 中使用的是简单的标量运算用于演示和学习自动微分的原理，但在实际的神经网络实现中，为了效率会进行大量的矩阵运算。

实现自动微分框架#

神经网络中有很多需要维护的参数，因此需要设计一个数据结构协助我们进行数据维护。

自定义类#

class Value:
    def __init__(self, data):
        self.data = data

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, other):
        out = Value(self.data + other.data)
        return out
    
    def __mul__(self, other):
        out = Value(self.data * other.data)
        return out

__repr__用于生成开发者友好的字符串表示
__str__用于生成用户有好的字符串表示通常由 print () 和 str () 调用
如果 str 不存在，print () 会退而使用 repr。
如果 repr 不存在，repr () 和交互式环境的直接显示会使用默认的 <...> 表示。

上面的代码中简单定义了一个Value类，包含属性 data，定义了 Value 类的方法 add 和 mul。

a = Value(2.0)
b = Value(-3.0)
c = Value(10.0)
d = a * b + c
d

Value(data=4.0)

接下来为 Value 增加一些其他的属性，使之能够保留更多的操作细节。这里使用集合_children表示当前 data 的子节点，op记录操作运算符，label 为当前数据的别称

class Value:
    def __init__(self, data, _children=(), _op='', label=''):
        self.data = data
        self._prev = set(_children)
        self._op = _op
        self.label = label

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, other):
        out = Value(self.data + other.data, (self, other), '+')
        return out
    
    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other), '*')
        return out

使用可视化 API 显示计算图#

为了更加直观的显示出计算过程，我们将通过下面的代码进行可视化操作（非重点，可跳过）

from graphviz import Digraph

def trace(root):
    nodes, edges = set(), set()
    def build(v):
        if v not in nodes:
            nodes.add(v)
            for child in v._prev:
                edges.add((child, v))
                build(child)
    build(root)
    return nodes, edges

def draw_dot(root):
    dot = Digraph(format='svg', graph_attr={'rankdir':'LR'})

    nodes, edges = trace(root)
    for n in nodes:
        uid = str(id(n))

        dot.node(name=uid, label='{%s | data %.4f | grad %.4f}' % (n.label, n.data, n.grad), shape='record')

        if n._op:
            dot.node(name=uid+n._op, label=n._op)
            dot.edge(uid + n._op, uid)
    
    for n1, n2 in edges:
        dot.edge(str(id(n1)), str(id(n2)) + n2._op)
    return dot

定义一个深层的表达式并可视化计算图（前向传播）

a = Value(2.0, label='a')
b = Value(-3.0, label='b')
c = Value(10.0, label='c')
e = a * b; e.label = 'e'
d = e + c; d.label = 'd'
f = Value(-2.0, label='f')
L = d * f; L.label = 'L'
L

draw_dot(L)

output

反向传播计算梯度#

在进反向传播之前先回顾一下链式求导法则

链式求导法则#

在反向传播过程中需要求解最终输出对中途某个输入的导数，但最终输出的函数并不是直接关于中途变量的函数。

eg. 如果 y 是关于 u 的函数 $(y=f(u))$ ，而 u 又是关于 x 的函数 $(u=g(x))$ ，则 $y$ 关于 $x$ 的导数为：

$\frac{dy}{dx} = \frac{dy}{du} \cdot \frac{du}{dx}$

如果 $h(x)$ 是复合函数 $f(g(x))$ ，那么 $h(x)$ 关于 $x$ 的导数 $h'(x)$ 等于外层函数 $f$ 的导数在内层函数 $g(x)$ 处的值，乘以内层函数 $g(x)$ 的导数 $g'(x)$ 。即：

$h'(x) = f'(g(x)) \cdot g'(x)$

求导时，先对外层函数求导（保持内层函数不变），然后乘以内层函数的导数。

手动反向求导#

修改 Value 类，实现手动的反向求导

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

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, 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):
        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

L 对 L 本身求导，导数为 1，所以先手动的将 L 的梯度设为 1

L.grad = 1.0

对 L 进行梯度反向传播，由计算图我们可以知道 L=f*d，那么 L 的反向传播就能分别求出 f 和 d 的梯度。像这样，手动的对 L 进行反向传播：

L._backward()

output
同理，对 d 进行反向传播得到 c、e 的梯度，对 c 反向传播得到 a 和 b 的梯度。

d._backward()
e._backward()

output

在神经网络中神经元接收数据后不仅进行简单的线性运算，还会对线性运算的结果进行非线性处理。我们接着拓展计算图，并使用激活函数 tanh。

定义 tanh 函数：

def tanh(self):
        n = self.data
        t = (math.exp(2*n) - 1) / (math.exp(2*n) + 1)
        out = Value(t, (self, ), 'tanh')
        def _backward():
            self.grad += (1 - t ** 2) * out.grad
        out._backward = _backward
        return out

tanh(x) = (e^2x - 1) / (e^2x + 1)
其导数为 tanh (x)' = 1 - tanh (x)²

定义一个复杂的神经元

x1 = Value(2.0, label='x1')
x2 = Value(0.0, label='x2')

w1 = Value(-3.0, label='w1')
w2 = Value(1.0, label='w2')

b = Value(6.8813735870195432, label='b')

x1w1 = x1 * w1; x1w1.label='x1w1'
x2w2 = x2 * w2; x2w2.label='x2w2'
x1w1x2w2 = x1w1 + x2w2; x1w1x2w2.label='x1w1x2w2'

n = x1w1x2w2 + b; n.label='n'
o = n.tanh(); o.label='o'

draw_dot(o)

output

像上面一样手动执行反向传播

output

如果计算过程再复杂一下，手动反向传播就过于繁杂了。接下来实现一个自动的反向传播。

自动反向传播#

具体来说，我们在前向传播后执行反向传播操作，按照拓扑排序的对节点的逆序调用 ._backward() 方法。

定义自动的反向传播函数

def backward(self):
        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)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

执行

o.backward()

即可一次性求出整个计算图的梯度。但是当前我们的设计有一定的缺陷，下面的例子能够清晰地显示出我们的缺陷：

a = Value(3.0, label='a')
b = a + a; b.label='b'
b.backward()
draw_dot(b)

b 对 a 的导数应该是 2，而不是 1，这是因为我们在定义反向传播时，只是简单的进行了赋值操作，对于多次使用某个变量的操作，梯度应该是累加的。

修复这个问题，完整的 Value 类为：

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

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, other):
      # 类型检验：允许Value类和数值进行运算，增强了Value类的鲁棒性和强健性
        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
    # __radd__是一个反向的魔术方法，使得处理2 * a这样的表达更自然
    def __radd__(self, other):
        return self.__add__(other)
    
    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
    
    def tanh(self):
        n = self.data
        t = (math.exp(2*n) - 1) / (math.exp(2*n) + 1)
        out = Value(t, (self, ), 'tanh')
        def _backward():
            self.grad += (1 - t ** 2) * out.grad
        out._backward = _backward
        return out 
    
    def backward(self):
        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)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

对于 Value 类中其他的相关操作，你可以自己再做定义，如 div 操作，pow 等。

视频中作者还讲解了一些其他的内容如使用 Pytorch 来实现一个简单的 MLP，感兴趣的话可以继续学习。

总结#

写下这篇博客，希望通过这趟 Micrograd 之旅，能让大家对自动微分不再感到神秘，甚至能自己动手实现一个 mini 的版本。如果这篇博客能帮助你对深度学习框架背后那些默默工作的 “英雄” 们（比如自动微分引擎）多一分理解，那我的目的就达到了！也欢迎大家继续探索，尝试给我们的 Value 类添加更多功能，或者去看看 Karpathy 原版的视频和代码，相信会有更多收获。