动手实现一个带自动微分的深度学习框架

动手实现一个带自动微分的深度学习框架

转自:Automatic Differentiation Tutorial

参考代码:https://github.com/borgwang/tinynn-autograd (主要看 core/tensor.py 和 core/ops.py)

动手实现一个带自动微分的深度学习框架_第1张图片

目录

  • 简介
  • 自动求导设计
  • 自动求导实现
  • 一个例子
  • 总结
  • 参考资料

简介

梯度下降(Gradient Descent)及其衍生算法是神经网络训练的基础,梯度下降本质上就是求解损失关于网络参数的梯度,不断计算这个梯度对网络参数进行更新。现代的神经网络框架都实现了自动求导的功能,只需要要定义好网络前向计算的逻辑,在运算时自动求导模块就会自动把梯度算好,不用自己手写求导梯度。

笔者在之前的 一篇文章 中讲解和实现了一个迷你的神经网络框架 tinynn,在 tinynn 中我们定义了网络层 layer 的概念,整个网络是由一层层的 layer 叠起来的(全连接层、卷积层、激活函数层、Pooling 层等等),如下图所示

动手实现一个带自动微分的深度学习框架_第2张图片

在实现的时候需要显示为每层定义好前向 forward 和反向 backward(梯度计算)的计算逻辑。从本质上看 这些 layer 其实是一组基础算子的组合,而这些基础算子(加减乘除、矩阵变换等等)的导函数本身都比较简单,如果能够将这些基础算子的导函数写好,同时把不同算子之间连接逻辑记录(计算依赖图)下来,那么这个时候就不再需要自己写反向了,只需要计算损失,然后从损失函数开始,让梯度自己用预先定义好的导函数,沿着计算图反向流动即可以得到参数的梯度,这个就是自动求导的核心思想。tinynn 中之所有 layer 这个概念,一方面是符合我们直觉上的理解,另一方面是为了在没有自动求导的情况下方便实现。有了自动求导,我们可以抛开 layer 这个概念,神经网络的训练可以抽象为定义好一个网络的计算图,然后让数据前向流动,让梯度自动反向流动( TensorFlow 这个名字起得相当有水准)。

我们可以看看 PyTorch 的一小段核心的训练代码(来源官方文档 MNIST 例子)

for batch_idx, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)
    optimizer.zero_grad()  # 初始化梯度
    output = model(data)  # 从 data 到 output 的计算图
    loss = F.nll_loss(output, target) # 从 output 到 loss 的计算图
    loss.backward()  # 梯度从 loss 开始反向流动
    optimizer.step()  # 使用梯度对参数更新

可以看到 PyTorch 的基本思路和我们上面描述的是一致的,定义好计算图 -> forward 得到损失 -> 梯度反向流动。

自动求导设计

知道了自动求导的基本流程之后,我们考虑如何来实现。先考虑没有自动求导,为每个运算手动写 backward 的情况,在这种情况下我们实际上定义了两个计算图,一个前向一个反向,考虑最简单的线性回归的运算 W X + B WX+B WX+B,其计如下所示。

动手实现一个带自动微分的深度学习框架_第3张图片

可以看到这两个计算图的结构实际上是一样的,只是在前向流动的是计算的中间结果,反向流动的是梯度,以及中间的运算反向的时候是导数运算。实际上我们可以把两者结合到一起,只定义一次前向计算图,让反向计算图自动生成

动手实现一个带自动微分的深度学习框架_第4张图片

从实现的角度看,如果我们不需要自动求导,那么网络框架中的 Tensor 类只需要对 Tensor 运算符有定义,能够进行数值运算(tinynn 中就简单的使用 ndarray 作为 Tensor 的实现)。但如果要实现自动求导,那么 Tensor 类需要额外做几件事:

  1. 增加一个梯度的变量保存当前 tensor 的梯度
  2. 保存当前 tensor 依赖的 tensor(如上图中 O1O1 依赖于 X,WX,W)
  3. 保存下对各个依赖 tensor 的导函数(这个导函数的作用是将当前 tensor 的梯度传到依赖的 tensor 上)

自动求导实现

我们按照上面的分析开始实现 Tensor 类如下,初始化方法中首先把 tensor 的值保存下来,然后有一个 requires_grad 的 bool 变量表明这个 tensor 是不是需要求梯度,还有一个 dependency 的列表用于保存该 tensor 依赖的 tensor 以及对于他们的导函数。

zero_grad() 方法比较简单,将当前 tensor 的梯度设置为 0,防止梯度的累加。自动求导从调用计算图的最后一个节点 tensor 的 backward() 方法开始(在神经网络中这个节点一般是 loss)。backward() 方法主要流程为

  • 确保改 tensor 确实需要求导 self.requires_grad == True
  • 将从上个 tensor 传进来的梯度加到自身梯度上,如果没有(反向求导的起点 tensor),则将梯度初始化为 1.0
  • 对每一个依赖的 tensor 运行保存下来的导函数,计算传播到依赖 tensor 的梯度,然后调用依赖 tensor 的 backward() 方法。可以看到这其实就是 Depth-First Search 计算图的节点
def as_tensor(obj):
    if not isinstance(obj, Tensor):
        obj = Tensor(obj)
    return obj


class Tensor:

    def __init__(self, values, requires_grad=False, dependency=None):
        self._values = np.array(values)
        self.shape = self.values.shape

        self.grad = None
        if requires_grad:
            self.zero_grad()
        self.requires_grad = requires_grad

        if dependency is None:
            dependency = []
        self.dependency = dependency

    @property
    def values(self):
        return self._values

    @values.setter
    def values(self, new_values):
        self._values = np.array(new_values)
        self.grad = None

    def zero_grad(self):
        self.grad = np.zeros(self.shape)

    def backward(self, grad=None):
        assert self.requires_grad, "Call backward() on a non-requires-grad tensor."
        grad = 1.0 if grad is None else grad
        grad = np.array(grad)

        # accumulate gradient
        self.grad += grad

        # propagate the gradient to its dependencies
        for dep in self.dependency:
            grad_for_dep = dep["grad_fn"](grad)
            dep["tensor"].backward(grad_for_dep)

可能看到这里读者可能会疑问,一个 tensor 依赖的 tensor 和对他们的导函数(也就是 dependency 里面的东西)从哪里来?似乎没有哪一个方法在做保存依赖这件事。

假设我们可能会这样使用我们的 Tensor 类

W = Tensor([[1], [3]], requires_grad=True)  # 2x1 tensor
X = Tensor([[1, 2], [3, 4], [5, 6], [7, 8]], requires_grad=True)  # 4x2 tensor
O = X @ W  # suppose to be a 4x1 tensor

如何让 XW 完成矩阵乘法输出正确的 O 的同时,让 O 能记下他依赖于 WX 呢?答案是重载运算符

class Tensor:
    # ...
    def __matmul__(self, other):
        # 1. calculate forward values
        values = self.values @ other.values

        # 2. if output tensor requires_grad
        requires_grad = ts1.requires_grad or ts2.requires_grad

        # 3. build dependency list
        dependency = []
        if self.requires_grad:
            # O = X @ W
            # D_O / D_X = grad @ W.T
            def grad_fn1(grad):
            	return grad @ other.values.T
            dependency.append(dict(tensor=self, grad_fn=grad_fn1))
        if other.requires_grad:
            # O = X @ W
            # D_O / D_W = X.T @ grad
			def grad_fn2(grad):
            	return self.values.T @ grad
            dependency.append(dict(tensor=other, grad_fn=grad_fn2))
        return Tensor(values, requires_grad, dependency)
    # ...

关于 Python 中如何重载运算符这里不展开,读者有兴趣可以参考官方文档或者这篇文章。基本上在 Tensor 类内定义了 __matmul__ 这个方法后,实际上是重载了矩阵乘法运算符 @ (Python 3.5 以上支持) 。当运行 X @ W 时会自动调用 X__matmul__ 方法。

这个方法里面做了三件事:

  1. 计算矩阵乘法结果(这个是必须的)

  2. 确定是否需要新生成的 tensor 是否需要梯度,这个由两个操作数决定。比如在这个例子中,如果 W 或者 X 需要梯度,那么生成的 O也是需要计算梯度的(这样才能够计算 W 或者 X 的梯度)

  3. 建立 tensor 的依赖列表

    自动求导中最关键的部分就是在这里了,还是以 O = X @ W 为例子,这里我们会先检查是否 X需要计算梯度,如果需要,我们需要把导函数 D_O / D_X 定义好,保存下来;同样的如果 W 需要梯度,我们将 D_O / D_W 定义好保存下来。最后生成一个 dependency 列表保存着在新生成的 tensor O 中。

然后我们再回顾前面讲的 backward()方法,backward() 方法会遍历 tensor 的 dependency ,将用保存的 grad_fn 计算要传给依赖 tensor 的梯度,然后调用依赖 tensor 的 backward() 方法将梯度传递下去,从而实现了梯度在整个计算图的流动。

grad_for_dep = dep["grad_fn"](grad)
dep["tensor"].backward(grad_for_dep)

自动求导讲到这里其实已经基本没有什么新东西,剩下的工作就是以类似的方法大量地重载各种各样的运算符,使其能够 cover 住大部分所需要的操作(基本上照着 NumPy 的接口都给重载一次就差不多了 )。无论你定义了多复杂的运算,只要重载了相关的运算符,就都能够自动求导了,再也不用自己写梯度了。

动手实现一个带自动微分的深度学习框架_第5张图片

一个例子

大量的重载运算符的工作在文章里就不贴上来了(过程不怎么有趣),我写在了一个 notebook 上,大家有兴趣可以去看看 borgwang/toys/ml-autograd。在这个 notebook 里面重载了实现一个简单的线性回归需要的几种运算符,以及一个线性回归的例子。这里把例子和结果贴上来

# training data
x = Tensor(np.random.normal(0, 1.0, (100, 3)))
coef = Tensor(np.random.randint(0, 10, (3,)))
y = x * coef - 3

params = {
    "w": Tensor(np.random.normal(0, 1.0, (3, 3)), requires_grad=True),
    "b": Tensor(np.random.normal(0, 1.0, 3), requires_grad=True)
}

learng_rate = 3e-4
loss_list = []
for e in range(101):
    # set gradient to zero
    for param in params.values():
        param.zero_grad()

    # forward
    predicted = x @ params["w"] + params["b"]
    err = predicted - y
    loss = (err * err).sum()

    # backward automatically
    loss.backward()

    # updata parameters (gradient descent)
    for param in params.values():
        param -= learng_rate * param.grad

    loss_list.append(loss.values)
    if e % 10 == 0:
        print("epoch-%i \tloss: %.4f" % (e, loss.values))

epoch-0 	loss: 8976.9821
epoch-10 	loss: 2747.4262
epoch-20 	loss: 871.4415
epoch-30 	loss: 284.9750
epoch-40 	loss: 95.7080
epoch-50 	loss: 32.9175
epoch-60 	loss: 11.5687
epoch-70 	loss: 4.1467
epoch-80 	loss: 1.5132
epoch-90 	loss: 0.5611
epoch-100 	loss: 0.2111

动手实现一个带自动微分的深度学习框架_第6张图片

接口和 PyTorch 相似,在每个循环里面首先将参数梯度设为 0 ,然后定义计算图,然后从 loss 开始反向传播,最后更新参数。从结果可以看到 loss 随着训练进行非常漂亮地下降,说明我们的自动求导按照我们的设想 work 了。

总结

本文实现了讨论了自动求导的设计思路和整个过程是怎么运作的。总结起来:自动求导就是在定义了一个有状态的计算图,该计算图上的节点不仅保存了节点的前向运算,还保存了反向计算所需的上下文信息。利用上下文信息,通过图遍历让梯度在图中流动,实现自动求节点梯度。

我们通过重载运算符实现了一个支持自动求导的 Tensor 类,用一个简单的线性回归 demo 测试了自动求导。当然这只是最基本的能实现自动求导功能的 demo,从实现的角度上看还有很多需要优化的地方(内存开销、运算速度等),笔者有空会继续深入研究,读者如果有兴趣也可以自行查阅相关资料。Peace out.

参考资料

  • PyTorch Doc
  • PyTorch Autograd Explained - In-depth Tutorial
  • joelgrus/autograd
  • Automatic Differentiation in Machine Learning: a Survey

你可能感兴趣的:(MLSys,深度学习,人工智能,神经网络)