PyTorch自动微分基本原理

Posted 2021-02-17 cocode

序言：在训练一个神经网络时，梯度的计算是一个关键的步骤，它为神经网络的优化提供了关键数据。但是在面临复杂神经网络的时候导数的计算就成为一个难题，要求人们解出复杂、高维的方程是不现实的。这就是自动微分出现的原因，当前最流行的深度学习框架如PyTorch、Tensorflow等都提供了自动微分的支持，让人们只需要很少的工作就能神奇般地自动计算出复杂函数的梯度。

PyTorch的autograd简介

Tensor是PyTorch实现多维数组计算和自动微分的关键数据结构。一方面，它类似于numpy的ndarray，用户可以对Tensor进行各种数学运算；另一方面，当设置.requires_grad = True之后，在其上进行的各种操作就会被记录下来，用于后续的梯度计算，其内部实现机制被成为动态计算图(dynamic computation graph)。

Variable变量：在PyTorch早期版本中，Tensor只负责多维数组的运算，自动微分的职责是Variable完成的，因此经常可以看到因而产生的包装代码。而在0.4.0版本之后，二者的功能进行了合并，使得自动微分的使用更加简单了。

autograd机制能够记录作用于Tensor上的所有操作，生成一个动态计算图。图的叶子节点是输入的数据，根节点是输出的结果。当在根节点调用.backward()的时候就会从根到叶应用链式法则计算梯度。默认情况下，只有.requires_grad和is_leaf两个属性都为True的节点才会被计算导数，并存储到grad中。

动态计算图本质上是一个有向无环图，因此“叶”和“根”的称呼是不太准确的，但是这种简称可以帮助理解，PyTorch的文档中仍然采用这种说法。

requires_grad属性

requires_grad属性默认为False，也就是Tensor变量默认是不需要求导的。如果一个节点的requires_grad是True，那么所有依赖它的节点requires_grad也会是True。换言之，如果一个节点依赖的所有节点都不需要求导，那么它的requires_grad也会是False。在反向传播的过程中，该节点所在的子图会被排除在外。

>>> x = torch.randn(5, 5)  # requires_grad=False by default
>>> y = torch.randn(5, 5)  # requires_grad=False by default
>>> z = torch.randn((5, 5), requires_grad=True)
>>> a = x + y
>>> a.requires_grad
False
>>> b = a + z
>>> b.requires_grad
True

Function类

我们已经知道PyTorch使用动态计算图(DAG)记录计算的全过程，那么DAG是怎样建立的呢？一些博客认为DAG的节点是Tensor(或说Variable)，这其实是不准确的。DAG的节点是Function对象，边表示数据依赖，从输出指向输入。因此Function类在PyTorch自动微分中位居核心地位，但是用户通常不会直接去使用，导致人们对Function类了解并不多。

每当对Tensor施加一个运算的时候，就会产生一个Function对象，它产生运算的结果，记录运算的发生，并且记录运算的输入。Tensor使用.grad_fn属性记录这个计算图的入口。反向传播过程中，autograd引擎会按照逆序，通过Function的backward依次计算梯度。

 

 

backward函数

backward函数是反向传播的入口点，在需要被求导的节点上调用backward函数会计算梯度值到相应的节点上。backward需要一个重要的参数grad_tensor，但如果节点只含有一个标量值，这个参数就可以省略（例如最普遍的loss.backward()与loss.backward(torch.tensor(1))等价），否则就会报如下的错误：

Backward should be called only on a scalar (i.e. 1-element tensor) or with gradient w.r.t. the variable

要理解这个参数的内涵首先要从数学角度认识梯度运算。如果有一个向量函数$vec{y}=f(vec{x})$，那么$vec{y}$相对于$vec{x}$的梯度是一个雅克比矩阵(Jacobian matrix)：

$$egin{split}J=left(egin{array}{ccc} frac{partial y_{1}}{partial x_{1}} & cdots & frac{partial y_{1}}{partial x_{n}}\\ vdots & ddots & vdots\\ frac{partial y_{m}}{partial x_{1}} & cdots & frac{partial y_{m}}{partial x_{n}} end{array} ight)end{split}$$

本文讨论的主角torch.autograd本质上是一个向量-雅克比乘积(*vector-Jacobian product*)的计算引擎，即计算$v^{T}cdot J$，而所谓的参数grad_tensor就是这里的$v$。由定义易知，参数grad_tensor需要与Tensor本身有相同的size。通过恰当地设置grad_tensor，容易计算任意的$frac{partial y_{m}}{partial x_{n}}$求导组合。

反向传播过程中一般用来传递上游传来的梯度，从而实现链式法则，简单的推导如下所示：

$$egin{split}J^{T}cdot v=left(egin{array}{ccc} frac{partial y_{1}}{partial x_{1}} & cdots & frac{partial y_{m}}{partial x_{1}}\\ vdots & ddots & vdots\\ frac{partial y_{1}}{partial x_{n}} & cdots & frac{partial y_{m}}{partial x_{n}} end{array} ight)left(egin{array}{c} frac{partial l}{partial y_{1}}\\ vdots\\ frac{partial l}{partial y_{m}} end{array} ight)=left(egin{array}{c} frac{partial l}{partial x_{1}}\\ vdots\\ frac{partial l}{partial x_{n}} end{array} ight)end{split}$$

(注：这里的计算结果被转置为列向量以方便查看)

---

注意：梯度是累加的

backward函数本身没有返回值，它计算出来的梯度存放在叶子节点的grad属性中。PyTorch文档中提到，如果grad属性不为空，新计算出来的梯度值会直接加到旧值上面。

为什么不直接覆盖旧的结果呢？这是因为有些Tensor可能有多个输出，那么就需要调用多个backward。叠加的处理方式使得backward不需要考虑之前有没有被计算过导数，只需要加上去就行了，这使得设计变得更简单。因此我们用户在反向传播之前，常常需要用zero_grad函数对导数手动清零，确保计算出来的是正确的结果。