理解 pytorch 中的雅可比张量梯度

Posted

技术标签:

【中文标题】理解 pytorch 中的雅可比张量梯度【英文标题】:Understanding Jacobian tensor gradients in pytorch 【发布时间】:2021-06-08 15:33:40 【问题描述】:

我正在浏览official pytorch tut,它解释了张量梯度和雅可比乘积如下:

PyTorch 允许您计算给定输入向量 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\n", inp.grad)
out.backward(torch.ones_like(inp), retain_graph=True)
print("\nSecond call\n", inp.grad)
inp.grad.zero_()
out.backward(torch.ones_like(inp), retain_graph=True)
print("\nCall after zeroing gradients\n", 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 输出是如何计算的。有人可以用这个例子中的计算来解释一下雅可比矩阵吗?

【问题讨论】:

【参考方案1】:

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

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

首先,我们应该看一下雅可比本身,它对应于包含所有偏导数J_ijmn = dy_mn/dx_ij 的张量J。根据y_mn 的表达式,我们可以说对于所有ijmndy_mn/dx_ij = d(x_mn + 1)²/dx_ij 如果m≠in≠j,则为0。否则,m=in=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

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

x.grad = dL/dx,形如xdL/dy 是传入梯度:backward 函数中的 gradient 参数 dL/dx 是上述的雅可比张量。

如文档中所述,应用 backward 实际上并不提供雅可比行列式。它直接计算链式法则积并存储梯度(dL/dxx.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) -&gt; 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.]])

这与从 outtorch.ones_like(inp) 作为传入梯度进行反向传播时得到的结果相匹配:

>>> 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,使您能够一次性处理和重置该集合中的所有参数。


我已将torch.autograd.functional 导入为A

【讨论】:

以上是关于理解 pytorch 中的雅可比张量梯度的主要内容,如果未能解决你的问题,请参考以下文章

PyTorch 线性代数梯度

一元函数的梯度和雅可比矩阵是否想用

Pytorch——张量 Tensors

Pytorch自动求解梯度

PyTorch进阶之路:张量与梯度

pytorch之求梯度和nn.Linear的理解