专栏名称: 小白学视觉

本公众号主要介绍机器视觉基础知识和新闻，以及在学习机器视觉时遇到的各种纠结和坑的心路历程。

深入浅出详解张量自动求导机制

小白学视觉 · 公众号 · · 2024-03-07 10:05

正文

点击上方 “ 小白学视觉 ”，选择加" 星标 "或“ 置顶 ”

重磅干货，第一时间送达

前言

自动求导机制是神经网络计算框架必备的组件，被赋予了多种称呼：Autograd、Autodiff、自动求导、自动梯度、自动微分。万变不离其宗，在神经网络的训练过程中，它被用来计算权重的最速下降方向，以指导优化器下一步迭代时对权重的更新。

作者 | 清川

转载自 | PaperWeekly

『写在前面』

深入浅出，在计算机教材界被用滥的词，总是继承着领域小白的初心和梦想。顾名思义，它既意味着理解得透彻，又要求复述得通俗。如果说复述是大名鼎鼎的费曼学习法的精髓，那么复刻便是其在程序世界最恰当的对应概念。君不见深度学习水涨船高，计算框架层出不穷。想要深入浅出这些框架，何不亲自动手复刻个轮子？

自动求导机制是神经网络计算框架必备的组件，被赋予了多种称呼：Autograd、Autodiff、自动求导、自动梯度、自动微分。万变不离其宗，在神经网络的训练过程中，它被用来计算权重的最速下降方向，以指导优化器下一步迭代时对权重的更新。有了自动求导（比如 JAX 框架的 grad），再辅之以矩阵运算（比如 Numpy），就可以实现神经网络的基本功能了。

最近和 @王桂波博主交流，受益良多。本文的框架使用 Python3 编写，主要参考了他的教程 Automatic Differentiation Tutorial [1] ，并补充了一些网上未曾讨论但很重要的细节。本文会持续更新，分析一些常见的算子设计，之后还会给出使用本文的框架做线性回归和神经网络训练的例子，敬请关注！

『反向传播』

这里以最简单的前馈神经网络为例（没有反馈链路的多层感知机。正向传播过程即是输入数据逐层推断，最后得到预测值的过程，是机器对已学知识的演练；反向传播则是比较预测值和真实值，根据定义的损失函数反向逐层归谬的过程，是自我批评、寻找不足之处的改进。

具体地，反向传播会求出损失值对于各权重的负梯度，来寻找改进的最佳方向。各层在推断时只是接收上层的信息、做决策、传递结果给下一层，所以在归谬时只需要各层自我反思即可。那么距离输出最远的输入层则归谬最为复杂。直觉上讲，它的决策经过了多层修改，想要判断好不好已经很模糊了。

『链式法则』

复合函数的链式法则实际上是以上反向逐层计算梯度的理论支持。以最简单的一元（标量）复合函数为例，链式法则如下：

类比神经网络，如果上式右侧连乘的每一项代表各层的局部梯度，则网络的输出对于各层的全局梯度就等于，从这一层开始到输出层的各层局部梯度的连乘。正因为输出层求梯度连乘只有一项，而输入层需要连乘所有项，求梯度的过程是反向逐层进行的。为什么可以这么类比呢？

一个两层的多层感知机可以定义如下，其中是输入，是各层输出，是权重矩阵，是各层偏置，是非线性激活函数：

根据上式，显然可以看做的一个复合函数：

只不过要注意，神经网络的复合函数是矩阵函数，而 矩阵函数的链式法则有自己的规律，并不是简单的点积！！！ 尽管没有统一的规律，要实现也并不困难，因为神经网络中能够用到的运算很有限，我们只需要按照矩阵粒度，将所有用到的运算的链式法则穷举定义即可。

以矩阵作为梯度的最小单位而不是神经网络的一层，是为了更灵活的表达能力：任何运算过程都可以按照计算顺序看做复合函数：

实现了广义的链式法则，就可以适配各种类型的网络定义了。

『总体框架』

4.1 封装张量类型

我们将具备自动求导功能的矩阵封装成一个叫做张量的类：

# import numpy as np
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

类的成员变量的作用如下：

values：通过初始化函数传入初值。被 property 装饰器定义为可读写的属性，主要为了在类外部对其进行赋值修改时控制其值始终为 Numpy 的 ndarray 类型。关于这部分 Python 语法可以参考 python 中的 property 装饰器 [2] ，是工程中常用的 Getter/Setter 设计模式 [3] 。

grad：存储该矩阵最终的全局梯度值。

requires_grad：表明该矩阵是否参与梯度计算。如果参与则给 grad 分配空间并初始化为与值形状相同的全零矩阵；如果不参与梯度计算，则不分配空间以优化内存效率。这里梯度清零操作定义如下：

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

dependency：当前矩阵可能储存的是某个运算的结果，我们需要记录其梯度如何向操作数矩阵传播。由于操作数可能不唯一，这个属性是列表类型。其中每一项将会是一个字典，字典的 tensor 字段指代操作数的张量，grad_fn 字段指代传播到该张量需要执行的函数。

4.2 实现反向传播

我们先来看一下反向传播的定义，稍后再讨论梯度清零功能的必要性：

class Tensor:
    # ...
    def




    
 backward(self, grad=None):
        assert self.requires_grad, "Call backward() on a non-requires-grad tensor."
        assert not (grad is None and self.values.size > 1), "grad can be implicitly created only for scalar outputs"
        grad = 1.0 if grad is None else grad
        grad = np.array(grad)

        self.grad += grad

        for dep in self.dependency:
            grad_for_dep = dep["grad_fn"](grad)
            dep["tensor"].backward(grad_for_dep)

在我们的程序逻辑里，当前张量最终的梯度是在上级函数处计算完毕后传进来的。仍以上面的算数运算复合函数为例：

张量的梯度就是由张量的反向传播函数计算好后传给的。这只是为了编程方便，没什么道理。如果当前张量就是计算图的最后一个节点（即反向传播开始的节点，比如），那么就不需要传入 grad，下面在判断到这种情况时会将 grad 初始化为 1.0。

函数中首先两个断言，第一个判断要求当前张量参与梯度计算，第二个判断要求当前张量为输出节点时必须是标量才能反向传播。在 Pytorch 中就不支持非标量对向量反向传播求梯度，否则会报以下错误：

RentimeError: grad can be implicitly created only for scalar outputs

这个很好理解：假如输出节点是向量，向量对向量求梯度会得到一个矩阵，而对矩阵求梯度就会得到更高维的矩阵，如此传播下去程序将会不可控；而标量对向量求梯度得到向量，对矩阵求梯度得到矩阵，程序可以链式传播。因此，我们必须在设计损失函数的时候就想办法压缩维度，让最后得到一个标量。

接下来将 现有的梯度上加上传入的梯度，这是一个有意思的 Trick。 在这些矩阵第一次参与反向传播，或被我们手动调用 zero_grad() 清零后，就相当于将张量的梯度直接赋值为传入的梯度。然而如果我们不进行梯度清零，这里的梯度就会累加。下面插播一下实现这种梯度累加的好处，同样也是梯度清零的意义。

BEGIN TIP

在神经网络训练中，我们使用的小批量梯度下降一般都是将一批样本的损失函数求和，然后进行反向传播。我们在定义网络的时候往往需要特别指定输入的形状，并且可以多留出一个维度用于设置弹性的 batchsize。现代计算框架一般都会提供使用 GPU 并行处理梯度计算图的能力，我猜想框架会隐式地为一批样本建立同样数量的计算图，同时加载到显存中。在计算好每个计算图节点的梯度后再进行累加合并。

这种猜想不无根据，之前在 Jetson TX2 上跑实验时发现 batchsize 超过 8 就会炸显存。TX2 是 CPU/GPU 共用 8G 内存的，而当时模型大小大概 260MB，使用的是 Adam 优化器。

我不知道有没有专门做模型训练内存/显存占用预测的论文，但这篇博客深度学习中 GPU 和显存分析 [4] 对显存占用讲得很透彻。据分析，Adam 优化器需要额外的 3 倍模型大小来存储梯度和动量等信息。那么假设模型的参数和梯度分别需要一份整体的备份在内存里，Adam 占用 3 倍模型大小，每次训练同时加载 8 份 Adam，则需要的总大小为：

再加上样本数据占用的内存以及操作系统占用的内存，确实跑起来很吃力。另一个例子来自论坛的问题 mxnet 有在有限显存的情况下增大 batchsize 的方法吗？ [5] ，题主使用 resnet152-v2 模型（根据 Keras documentation: Keras Applications [6] 、精度为单精度浮点数计算，模型大小约 230MB。）在显存大小为 6GB 的 GTX1060 上也只能把 batchsize 设成 16。

参考 PyTorch 中在反向传播前为什么要手动将梯度清零？ [7] ，其实我们可以逐个样本进行梯度计算，只要不在每轮都进行梯度清零，梯度就会一直累加，和批处理的效果是近似的。比如每隔 8 个样本进行一次梯度清零和优化器的 step，就相当于 batchsize 设成了 8。这样我们可以在有限的显存条件下尽可能加大 batchsize。当然这么做也是有代价的，本身并行的梯度计算变成了串行的，更额外引入了多次访存的开销，程序会变得很慢。

END TIP

继续分析反向传播的代码，下面这部分可谓是反向传播的精髓：

# ...
for dep in self.dependency:
    grad_for_dep = dep["grad_fn"](grad)
    dep["tensor"].backward(grad_for_dep)

循环遍历当前张量（操作结果）对应的操作数，调用 grad_fn 指向的梯度传播函数，将下一个节点梯度计算出来，递归调用该节点的反向传播函数，同时传入计算好的梯度。总体上看，这是一个对计算树深度优先遍历的过程。之前我们一直说的是计算图，但是想一想常见运算一般都是单/多输入单输出，多个输出的情况很少（比如 divmod，从输出节点倒推就是树形结构。

4.3 框架设计小结

整理一下，自动求导的整体框架如下图所示，这里使用的例子是：

▲ Autograd 总体框架：绘图不易，转载请注明出处

可以看到连接线密集的部分就是 mul 这个算子，这是反向传播实现的核心部分，接下来我们具体展开，看看算子如何定义，其梯度传播的 grad_fn 究竟长什么样子。

『算子设计』

我们把张量支持的运算函数叫做算子。为了更符合用户的使用习惯，以及表达式定义更自然，算子往往通过张量的运算符重载实现。算子函数内部既要实现正向传播时正常的运算功能，又要提供反向传播时梯度计算和传递的规则。下面我以三种最常见类型的运算为例，介绍一下算子的设计方法。

5.1 矩阵乘法算子

我们以矩阵乘法为例。在 Python 的 Numpy 库中使用 @ 运算符表示矩阵乘法，对应的运算符重载函数为 __matmul__。

def as_tensor(obj):
    if not isinstance(obj, Tensor):
        obj = Tensor(obj)
    return obj

class Tensor:
    # ...
    def __matmul__(self, other):
        # 0. make sure other is Tensor
        other = as_tensor(other)

        # 1. calculate forward values
        values = self.values @ other.values

        # 2. if output tensor requires_grad
        requires_grad = self.requires_grad or other.requires_grad

        # 3. build dependency list
        dependency = []
        if self.requires_grad:
            def grad_fn1(grad):
                pass  # TODO HERE
            dependency.append(dict(tensor=self, grad_fn=grad_fn1))
        if other.requires_grad:
            def grad_fn2(grad):
                pass  # TODO HERE
            dependency.append(dict(tensor=other, grad_fn=grad_fn2))
        return Tensor(values, requires_grad, dependency)

首先要保证另一个操作数是张量，在矩阵乘法中一般不会出现问题，但在数乘中，other 可能就只是一个数字。然后 values 直接计算矩阵乘法结果，作为返回的张量的值。之后，两个操作数只要有一个需要计算张量，结果就需要计算张量，否则计算树就截断了。最后构造依赖项，我们需要分别定义操作结果对两个操作数的梯度求解和传递过程。如果是标量乘法，这里毫无疑问会很简单：

在的乘法运算符重载函数被调用时，只需要将对和的梯度计算和传递分别定义如下：

在从开始反向传播，传播到的时候只需要把传入 grad_fn 函数即可。但如前文提到的，矩阵的链式法则有自己的规律，我们需要推导一下。这里的推导过程参考了一篇我十分佩服的文章：长躯鬼侠：矩阵求导术 [8] 。我们先从多元标量函数入手分析，根据标量的全微分公式、微分与梯度的关系，有：

那么类比标量，多元矩阵函数也应有如下关系：

其中 tr 代表矩阵的迹，即对角线元素之和。上式很好验证，我们假设，将上式右侧展开，很明显只有对角线元素的梯度和微分是对应的。

接下来，我们假设如下情形：

其中是的矩阵，是的矩阵，是的矩阵，是标量。可以理解为是个输入样本，每个样本有维特征，经过线性变换矩阵变换成高维特征图，然后通过（非线性变换以及损失函数等）计算得到最终的。下面计算对于的梯度，首先求的全微分：

考虑对求梯度时，可被看做常数，那么：

代入到全微分公式中，得到：

所以我们推导出对求梯度的链式法则：

同理，我们也能推导出对求梯度的链式法则：

于是，我们可以将上述代码中的 grad_fn 补全：

def grad_fn_1(grad):
    return np_matmul(grad, other.values.T)

def grad_fn_2(grad):
    return np_matmul(self.values.T, grad)

为了编程方便，我们将算子设计的代码中通用的部分封装起来：

def build_binary_ops(this, that, values, grad_fn_1, grad_fn_2):
    requires_grad = this.requires_grad or that.requires_grad
    dependency = []
    if this.requires_grad:
        dependency.append(dict(tensor=this, grad_fn=grad_fn_1))
    if that.requires_grad:
        dependency.append(dict(tensor=that, grad_fn=grad_fn_2))
    return this.__class__(values, requires_grad, dependency)

5.2 求平均值算子

在上述矩阵乘法的例子中 f 将最终的输出从矩阵映射成为标量，这个过程中往往先使用损失函数将矩阵变为向量，例如交叉熵损失函数：

之后我们需要再对一批样本的损失值向量做维度压缩，比如使用求平均值算子：

由于上文提到输出节点必须为标量，所以这一类函数也会非常常用。

class Tensor




    
:
    # ...
    def reduce_mean(self, axis=None):
        values = self.values.mean(axis=axis)
        def grad_fn(grad):
            grad = grad / self.values.size * np.ones_like(self.values)
            return grad
        return build_unary_ops(self, values, grad_fn)

首先也是实现正向传播的功能，直接调用 Numpy 数组自带的平均值函数 mean，得到输出。由于输出是一个标量，意味着上一层传来的 grad 也是个标量，那么这里反向传播的梯度计算和传递很简单，求和后的输出对于原始向量的每个元素的偏导数都为 1/n，所以只需要新建一个与 grad 相同维度的数组，然后再通过数乘进行 broadcast 即可。

但是问题不是这么简单，像求和、求均值这类运算往往提供了一个额外的参数：运算的轴。当操作的张量高于一维，我们就需要沿着轴去分配偏导的值。直接讲分配不好理解，举个例子：

其中，矩阵沿着 axis=1 求均值的操作可以写作：

而输出对于求导是。显然对矩阵求梯度最终的结果应为：

其实均值算子这里的梯度求解和传递就是：将上层传来的梯度，先沿着均值的轴扩展一个维度，然后再沿着这个轴进行重复，最终除以该轴的元素数量。在求均值时，比如一个 2×3×4 的矩阵，沿着 1 轴求均值，结果的形状就变成 2×4，也就是沿哪轴求值，哪轴就被压缩掉。

所以在反向传播时反其道，先使用 Numpy 的 expand_dims 将这个轴扩充。扩充后这个轴的长度只有 1，那么要扩充为原来的形状，就需要重复，重复多少次呢，当然重复原来这个轴上的元素数量那么多次。

有人可能疑惑，这里不也是向量对矩阵求梯度吗，为什么没出现高维矩阵？那是因为沿轴操作很特殊，是介于逐元素操作和矩阵操作之间的，不能按照传统矩阵乘法那类操作来类比。最后我们补全代码：

class Tensor:
    # ...
    def reduce_mean




    
(self, axis=None):
        values = self.values.mean(axis=axis)

        if axis is not None:
            repeat = self.values.shape[axis]

        def grad_fn(grad):
            if axis is None:
                grad = grad / self.values.size * np.ones_like(self.values)
            else:
                grad = np.expand_dims(grad / repeat, axis)
                grad = np.repeat(grad, repeat, axis)
            return grad

        return build_unary_ops(self, values, grad_fn)

5.3 broadcast算子

举一个最常见的例子，就是神经网络线性层中的偏置值。我们假设输入为，维度为，即个样本，每个样本维特征。设权重矩阵为，为维，表示每个特征的重要性。设偏置值为，激活函数为。那么线性层的前向传播过程应为：

这里