深扒torch.autograd原理

Posted 2022-12-19

当我们训练神经网络的时候，最常用的就是使用反向传播。在反向传播过程中，参数变化只和给定参数通过loss函数计算的梯度（gradient）有关。

PyTorch的torch.autograd提供了自动梯度计算，可以用于自动计算任何计算图的梯度。

举个简单的例子嗷，假设我们只有一层神经网络。输入为$x$，权重是$w$，bias是$b$，这里使用二元交叉熵（binary_cross_entropy）损失进行计算。

我们直接上计算图，图是pytorch官网的一张图：

这里$z$往前这一块是单层神经网络的计算，$y$是你的ground truth，用$z$和$y$计算二元交叉熵损失。

上面的整个过程可以用如下代码实现。

• 写法一：

  import torch
  
  x = torch.ones(5)  # input tensor
  y = torch.zeros(3)  # expected output
  w = torch.randn(5, 3, requires_grad=True)
  b = torch.randn(3, requires_grad=True)
  z = torch.matmul(x, w)+b
  loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)
  

• 写法二：

  import torch
  
  x = torch.ones(5)  # input tensor
  y = torch.zeros(3)  # expected output
  w = torch.randn(5, 3)
  b = torch.randn(3)
  z = torch.matmul(x, w)+b
  loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)
  w.requires_grad_(True)
  b.requires_grad_(True)
  

$w$和$b$是我们的可学习参数，所以这里我们要手动将其设置一下，告诉pytorch说我们需要计算这两个参数的梯度。

两个写法的区别就是你要在声明变量的时候直接使用requires_grad = True设置还是声明之后单独设置x.requires_grad_(True)。

我们使用Function对象构造计算图，它可以计算forward过程，并直接计算器在反向传播过程中的导数，反向传播过程存储在属性grad_fn中。

print(f"Gradient function for z = z.grad_fn")
print(f"Gradient function for loss = loss.grad_fn")

输出如下:

计算梯度

我们让神经网络学习的目的是想要优化网络的参数，也就是在学习过程中我们要使用损失函数计算参数的导数。我们要计算的有$ \\frac\\partial loss\\partial w$和$\\frac\\partial loss\\partial b$。

这一步需要我们手动调用loss.backward()进行计算，之后我们就可以使用w.grad和b.grad查看两个参数的导数值。

loss.backward()
print(w.grad)
print(b.grad)

这里有几点注意事项：

• 我们只能使用grad来求计算图中设置了requires_grad = True的叶子节点的梯度，没办法获取其他节点的梯度。

• 我们对一个计算图使用loss.backward()的时候只能用一次，但是为了模型的性能，我们可能要多次对同一个图进行loss.backward()，我们需要在调用loss.backward()的时候传入参数retain_graph=True。

关闭梯度追踪

默认情况下，当你设置了requires_grad = True之后，会自动记录该变量的计算历史，并将其用到梯度计算中，但是有的情况下我们并不想要计算历史数据。比如，我们把模型训练好了以后，我们只想直接用，把输入丢给它，让它计算结果，而不再需要它继续学习进行反向传播了，在这种情况下我们就要使用torch.no_grad()停止计算梯度追踪。

z = torch.matmul(x, w)+b
print(z.requires_grad)

with torch.no_grad():
    z = torch.matmul(x, w)+b
print(z.requires_grad)

还有另一种方法，不过不太常见，使用detach()，写法如下：

z = torch.matmul(x, w)+b
z_det = z.detach()
print(z_det.requires_grad)

关闭梯度追踪的情况：

• 想冻结住你模型的一部分，将其中一些参数变为冻结参数。比如你微调预训练模型的时候。

• 加速计算，当你只想计算前向计算过程的时候，就是我上边提到的情况，不使用梯度追踪可以让计算更高效。

---

关于autograd的更多的知识

还是看这张图，从概念上讲，autograd在一个由Function对象组成的有向无环图(DAG)中保存数据(张量)和所有执行的操作(以及产生的新张量)的记录。

在这个有向无环图中，叶子节点是是输入张量，根节点是输出张量。通过从根到叶追踪这个图的导数计算，可以使用链式法则自动计算梯度。

forward过程中，autograd同时在做两件事：

• 按照要求的操作计算，并获得最终的结果张量

• 维持有向无环图中计算的梯度函数

backward过程中，autograd做的是：

• 从之前保存的.grad_fn中计算梯度

• 将结果累积存储在张量的.grad属性中

• 使用链式法则进行反向传播，直到叶子节点。