理解火炬中的雅可比张量梯度

问题描述

我正在学习official pytorch tut，其中对张量梯度和雅可比乘积的解释如下：

不是计算雅可比矩阵本身，而是允许计算给定输入向量v=(v1…)的雅可比乘积Vm)。这是通过使用v作为参数回调来实现的：

inp = torch.eye(5, requires_grad=True)
out = (inp+1).pow(2)
out.backward(torch.ones_like(inp), retain_graph=True)
print("First call
", inp.grad)
out.backward(torch.ones_like(inp), retain_graph=True)
print("
Second call
", inp.grad)
inp.grad.zero_()
out.backward(torch.ones_like(inp), retain_graph=True)
print("
Call after zeroing gradients
", inp.grad)

结果：

First call
 tensor([[4., 2., 2., 2., 2.],
        [2., 4., 2., 2., 2.],
        [2., 2., 4., 2., 2.],
        [2., 2., 2., 4., 2.],
        [2., 2., 2., 2., 4.]])

Second call
 tensor([[8., 4., 4., 4., 4.],
        [4., 8., 4., 4., 4.],
        [4., 4., 8., 4., 4.],
        [4., 4., 4., 8., 4.],
        [4., 4., 4., 4., 8.]])

Call after zeroing gradients
 tensor([[4., 2., 2., 2., 2.],
        [2., 4., 2., 2., 2.],
        [2., 2., 4., 2., 2.],
        [2., 2., 2., 4., 2.],
        [2., 2., 2., 2., 4.]])

虽然我知道什么是雅可比矩阵，但我不知道这个雅可比乘积是如何计算的。

以下是我试图打印出来以获得理解的不同张量：

>>> out
tensor([[4., 1., 1., 1., 1.],
        [1., 4., 1., 1., 1.],
        [1., 1., 4., 1., 1.],
        [1., 1., 1., 4., 1.],
        [1., 1., 1., 1., 4.]], grad_fn=<PowBackward0>)
>>> torch.eye(5)
tensor([[1., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0.],
        [0., 0., 0., 1., 0.],
        [0., 0., 0., 0., 1.]])
>>> torch.ones_like(inp)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])
>>> inp
tensor([[1., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0.],
        [0., 0., 0., 1., 0.],
        [0., 0., 0., 0., 1.]], requires_grad=True)

但我不知道TUTS的产量是如何计算出来的。谁能用这个例子中的计算来解释一点雅可比矩阵？

---

解决方案

我们将经历整个过程：从计算雅可比到应用它以获得此输入的结果梯度。我们正在查看操作f(x) = (x + 1)²，在简单的标量设置中，我们得到df/dx = 2(x + 1)作为完全导数。

在多维设置中，我们有一个输入x_ij和一个输出y_mn，分别由(i, j)和(m, n)索引。函数映射定义为y_mn = (x_mn + 1)²。

首先，我们应该看看雅可比本身，这对应于包含所有偏导数的张量JJ_ijmn = dy_mn/dx_ij。从y_mn的表达式可以看出，对于所有i、j、m和n：dy_mn/dx_ij = d(x_mn + 1)²/dx_ij，0IFm≠i或n≠j。否则，即m=i或n=j，我们有d(x_mn + 1)²/dx_ij = d(x_ij + 1)²/dx_ij = 2(x_ij + 1)。

因此，J_ijmn可以简单定义为

         ↱ 2(x_ij + 1) if i=m, j=n
J_ijmn = 
         ↳ 0 else

从规则链中，输出相对于输入的梯度表示为dL/dx = dL/dy*dy/dx。从PyTorch的角度来看，我们有以下关系：

• x.grad = dL/dx，形状如x，
• dL/dy是传入的渐变：backward函数中的gradient参数
• dL/dx是上述雅可比张量。

如文档中所述，应用backward实际上并不提供雅可比。它直接计算链式规则乘积并存储梯度(，即dL/dx内部x.grad)。

就形状而言，雅可比乘法dL/dy*dy/dx = gradient*J将自身缩减为与x相同形状的张量。

执行的操作定义为：[dL/dx]_ij = ∑_mn([dL/dy]_ij * J_ijmn)。

---

如果我们将此应用于您的示例。我们有x = 1(i=j)(其中1(k): (k == True) -> 1是indicator function)，本质上只是单位矩阵。

我们计算雅可比矩阵：

         ↱ 2(1(i=j) + 1) =  if i=m, j=n
J_ijmn = 
         ↳ 0 else

变为

         ↱ 2(1 + 1) = 4  if i=j=m=n
J_ijmn = → 2(0 + 1) = 2  if i=m, j=n, i≠j
         ↳ 0 else

---

出于可视化目的，我们将继续使用x = torch.eye(2)：

>>> f = lambda x: (x+1)**2
>>> J = A.jacobian(f, inp)
tensor([[[[4., 0.],
          [0., 0.]],

         [[0., 2.],
          [0., 0.]]],


        [[[0., 0.],
          [2., 0.]],

         [[0., 0.],
          [0., 4.]]]])

然后使用torch.einsum计算矩阵乘法(我不会详细介绍，请查看this，然后this以深入了解EinSum求和运算符)：

>>> torch.einsum('ij,ijmn->mn', torch.ones_like(inp), J)
tensor([[4., 2.],
        [2., 4.]])

这与使用torch.ones_like(inp)作为传入渐变从out反向传播时获得的结果相匹配：

>>> out = f(inp)
>>> out.backward(torch.ones_like(inp))
>>> inp.grad
tensor([[4., 2.],
        [2., 4.]])

如果向后传播两次(当然，同时保留该图)，您最终将计算累加在参数的grad属性上的相同操作。因此，自然地，在两次向后传递之后，您将获得两次梯度：

>>> out = f(inp)
>>> out.backward(torch.ones_like(inp), retain_graph=True)
>>> out.backward(torch.ones_like(inp))
>>> inp.grad
tensor([[8., 4.],
        [4., 8.]])

这些渐变会累积，您可以通过调用inplace函数zero_：inp.grad.zero_()来重置它们。从那里开始，如果您再次反向传播，您将只保留一个累计渐变。

在实践中，您将在optimizer上注册参数，您可以从zero_grad调用zero_grad，从而一次性处理和重置该集合中的所有参数。

---

_{我已将torch.autograd.functional导入为A}