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

2022-12-24 13:12

©作者 | 清川

单位 | 上海交通大学博士生

研究方向 | 联邦学习、端云协同推断

写在前面

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

自动求导机制是神经网络计算框架必备的组件，被赋予了多种称呼：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算子

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

这里的结果为维，如果不支持算子广播，那么就要求也必须是维。然而在神经网络训练时往往等同于 batchsize，其大小是用户设置的，况且对于偏置值，创建倍的空间存储相同的值也是低效的。出于对空间效率和开发便捷性的考虑，我们就引入了广播机制。仅创建形状为 (1,) 的偏置张量，在相加时让这种加法操作广播到结果的每一个元素上。

上例中引入的支持广播的加法就是典型的 broadcast 算子（此外还包括数乘），概括地定义一下，就是把算子间形状相合的部分进行计算，形状不足的部分进行广播，从而降低对操作数形状上的要求，使用起来更加便捷。直接定义很难想象，现在的矩阵运算库基本都支持了算子的广播机制，我们以 Numpy 数组的加法为例：

x = np.zeros((2, 3, 4))
--------------------------------
array([[[0., 0., 0., 0.],  # 3x4
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],  # 3x4
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])

先初始化一个三维全零数组，为了展示运算作用的维度，我们在各个维度选择不同的元素数量。我们分别用形状为 (1,)、(4,)、(3, 1)、(3, 4)、(2, 3, 4)、(1, 3, 1)、(2,) 的随机数向量与 x 相加，为了控制结果的可复现性，我们使用固定随机数种子的 RandomState 产生向量。

# (1,)向量或者标量与x相加，将加在每一个元素上
>>> x +  np.random.RandomState(0).rand(1,)
array([[[0.55, 0.55, 0.55, 0.55],
        [0.55, 0.55, 0.55, 0.55],
        [0.55, 0.55, 0.55, 0.55]],

       [[0.55, 0.55, 0.55, 0.55],
        [0.55, 0.55, 0.55, 0.55],
        [0.55, 0.55, 0.55, 0.55]]])

# (4,)向量x相加，将x的2x3个形状为(4,)的子数组与该向量对应位置元素相加
>>> x +  np.random.RandomState(0).rand(4,)
array([[[0.55, 0.72, 0.6 , 0.54],
        [0.55, 0.72, 0.6 , 0.54],
        [0.55, 0.72, 0.6 , 0.54]],

       [[0.55, 0.72, 0.6 , 0.54],
        [0.55, 0.72, 0.6 , 0.54],
        [0.55, 0.72, 0.6 , 0.54]]])

# (3, 1)向量与x相加，沿着1轴对应位置元素相加，沿其他轴重复
>>> x +  np.random.RandomState(0).rand(3, 1)
array([[[0.55, 0.55, 0.55, 0.55],
        [0.72, 0.72, 0.72, 0.72],
        [0.6 , 0.6 , 0.6 , 0.6 ]],

       [[0.55, 0.55, 0.55, 0.55],
        [0.72, 0.72, 0.72, 0.72],
        [0.6 , 0.6 , 0.6 , 0.6 ]]])

# (3, 4)向量与x相加，将x的每一组3x4的子数组与该向量对应元素相加
>>> x +  np.random.RandomState(0).rand(3, 4)
array([[[0.55, 0.72, 0.6 , 0.54],
        [0.42, 0.65, 0.44, 0.89],
        [0.96, 0.38, 0.79, 0.53]],

       [[0.55, 0.72, 0.6 , 0.54],
        [0.42, 0.65, 0.44, 0.89],
        [0.96, 0.38, 0.79, 0.53]]])

# (2, 3, 4)向量与x相加，对应位置元素相加
>>> x +  np.random.RandomState(0).rand(2, 3, 4)
array([[[0.55, 0.72, 0.6 , 0.54],
        [0.42, 0.65, 0.44, 0.89],
        [0.96, 0.38, 0.79, 0.53]],

       [[0.57, 0.93, 0.07, 0.09],
        [0.02, 0.83, 0.78, 0.87],
        [0.98, 0.8 , 0.46, 0.78]]])

# (2,)向量与x相加，将报维度不匹配的错误
>>> x +  np.random.RandomState(0).rand(2,)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: operands could not be broadcast together with shapes (2,3,4) (2,)

由此我们可以总结，对于加法这类 broadcast 算子，并不像矩阵乘法要求乘数的形状完全一致，只要求至少有一个轴上的形状一致，其他轴为空或者为 1，即可进行运算。

以 x 与形状为（2, 3, 1）的向量 y 的加法为例，由于 x 的形状为（2, 3, 4），两个操作数 0 轴和 1 轴形状一致，而 y 缺少 2 轴（2 轴形状为 1，常常被认为是冗余的，可以 squeeze 掉的维度），那么操作的结果就是 0 轴和 1 轴形成的子矩阵的元素对应位置相加。

你可以把这个子矩阵的每个元素看成是一个向量，x、y 的子矩阵的元素分别是长度为 4 和 1 的向量，那么在对应位置元素相加时，又递归地发生了形状分别为 (1,) 和 (4,) 的向量的加法 broadcast。当然，你也可以只取 2 轴的第一个切片，认为加法操作是作用在 0 轴和 1 轴形成的子矩阵上（这时每个元素就是一个数值），但这种操作沿着 2 轴进行重复，总共重复了 4 次（2 轴的长度）。

接下来我们看看加法算子的逻辑具体怎么编写。前向传播直接使用加号就可以，因为我们张量使用的内部存储类型是 Numpy 的 ndarray，天然支持广播。

def __add__(ts1, ts2):
    ts2 = as_tensor(ts2)
    values = ts1.values + ts2.values
    # ...

重点是反向传播的逻辑，即 grad_fn 函数如何编写。在编写之前，我们要确定这种加法广播的结果如何对操作数进行求导。首先，将操作数全部转换为传统加法所要求的形状一致的情形。

假如加数 x、y 的形状分别为（2, 3, 4）和（3, 4），则将 y 的形状先扩充为（2, 3, 4），扩充的方法即按照形状不匹配的轴（0轴）重复 2 次（0 轴的长度），那么在相加时就相当于 y 广播到了其他的子矩阵上。这里「重复」的定义与 Numpy 的 expand_dims、repeat 函数的作用相同，前文在介绍平均值算子的时候提到过：

>>> y
array([[0.55, 0.72, 0.6 , 0.54],
       [0.42, 0.65, 0.44, 0.89],
       [0.96, 0.38, 0.79, 0.53]])
>>> y = np.expand_dims(y, 0)
>>> y
array([[[0.55, 0.72, 0.6 , 0.54],
        [0.42, 0.65, 0.44, 0.89],
        [0.96, 0.38, 0.79, 0.53]]])
>>> y = np.repeat(y, 2, axis=0)
>>> y
array([[[0.55, 0.72, 0.6 , 0.54],
        [0.42, 0.65, 0.44, 0.89],
        [0.96, 0.38, 0.79, 0.53]],

       [[0.55, 0.72, 0.6 , 0.54],
        [0.42, 0.65, 0.44, 0.89],
        [0.96, 0.38, 0.79, 0.53]]])

其实这种相似性很微妙，即加法的正向传播是维度填充，和 sum() 的反向传播相似；那么 sum() 的正向传播是维度缩减，就应该和加法的反向传播相似。可以想象，在求导后对应加数的局部导数必为形状与加数相同的全 1 矩阵，而加数扩充为（2, 3, 4）是我们假设的，其真实形状为（3, 4），那么每个元素的权重就应该变为 2 倍。

实际上我们反向传播时操作的是上一节点传来的全局导数，它在各个轴上的元素可能都不相等，这里就不是简单的乘法了，而是沿形状不匹配的轴（0轴）进行 sum reduce。

那么我们总结，grad_fn 中就是对缺轴的操作数求导时，要将上层传来的全局导数沿着形状不匹配的轴进行求和。

def grad_fn_ts1(grad):
    # handle broadcasting (5, 3) + (3,) -> (5, 3)
    for _ in range(grad.ndim - ts1.values.ndim):
        grad = grad.sum(axis=0)
    # handle broadcasting (5, 3) + (1, 3) -> (5, 3)
    for i, dim in enumerate(ts1.shape):
        if dim == 1:
            grad = grad.sum(axis=i, keepdims=True)
    return grad

注意，对于轴的长度为 1 的情况，我们要做 sum，但不要 reduce，使用 keepdims=True 来保持该轴不被 squeeze 掉。对右加数 ts2 的反向传播完全类似，数乘与加法完全类似，大家可以举一反三，也可以参考文末给出的代码。

5.4 用户自定义算子

为了良好的扩展性，很多神经网络计算框架，例如 TensorFlow，都支持用户自定义算子。我们的框架想要扩展也很简单，只需要实现以下函数原型即可：

# 一元运算
def unary_operation(operand, *args, **kwargs):

    # forward
    values = unary_operation_forward(operand)

    # backward
    def grad_fn(grad):
        # grad = ...
        return grad

    return build_unary_ops(operand, values, grad_fn)

# 二元运算（多元运算以此类推）
def binary_operation(operand_1, operand_2, *args, **kwargs):

    # forward
    values = binary_operation_forward(operand_1, operand_2)

    # backward
    def grad_fn_1(grad):
        # grad = ...
        return grad
    def grad_fn_2(grad):
        # grad = ...
        return grad

    return build_binary_ops(
        operand_1, operand_2, values, grad_fn_1, grad_fn_2)

注意事项

在 Numpy 中有个很讨厌的机制：一维数组无法转置。我们的框架内部一直使用的是数组来存放 values 和 grad，如果在梯度传递时出现了列向量与行向量做矩阵乘法的情况，本应得到矩阵，最终只会得到一个内积的标量。如果我们全部使用 matrix 来存储呢？则在很多情况下会出现多余的维度，需要不停的 squeeze。最终我只好对 @ 操作做了一层包装：

def np_matmul(arr1, arr2):
    if arr1.ndim == 1 and arr2.ndim == 1:
        arr1 = np.mat(arr1).T
        arr2 = np.mat(arr2)
    return arr1 @ arr2

这个问题是从 Automatic Differentiation Tutorial [9] 这篇文章给出的代码中发现的，大家有兴趣想复现的可以尝试一下。

总结

本文首先介绍了深度学习中常用的自动求导机制的原理和实现方法：将矩阵封装成张量并重载运算符，在正向传播表达式定义的同时，将反向传播的梯度计算和传递函数注册在操作结果的 dependency 表中，然后从输出节点反向深度优先遍历计算树，最后将计算好的全局梯度存储在张量的 grad 中。本文虽长，但仍无法做到面面俱到。希望大家能有所收获，反正我在写这篇文章的时候收获颇多。欢迎来讨论~