Pytorch 对输入而不是输出求和雅可比答案

【问题标题】：Pytorch sum jacobian over inputs instead of outputsPytorch 对输入而不是输出求和雅可比
【发布时间】：2021-10-29 09:05:42
【问题描述】：

假设我有一个张量 Y，它（直接或间接）从张量 X 计算而来。

通常当我应用torch.autograd.grad(Y, X, grad_outputs=torch.ones_like(Y)) 时，我会得到一个与X 形状相同的渐变蒙版。这个掩码实际上是Y w.r.t 的元素梯度的加权和。 X.

是否可以得到一个与Y 形状相同的渐变蒙版，其中每个元素mask[i][j] 是Y[i][j] w.r.t 的渐变之和。 X?

这相当于在 X 的维度上对雅可比 J(Y,X) 求和，而不是在 Y 的维度上求和。

>>> X = torch.eye(2)
>>> X.requires_grad_()
# X = [1 0]
#     [0 1]

>>> Y = torch.sum(X*X, dim=0)
# Y = [1, 1]

>>> torch.autograd.grad(Y, X, grad_outputs=torch.ones_like(Y), retain_graph=True)
(tensor([[2., 0.],
         [0., 2.]]),)

但是，我想要：

# [2, 2]

因为torch.sum(torch.autograd.grad(Y[0],X) 等于2 并且torch.sum(torch.autograd.grad(Y[1],X) 也等于2。

很容易计算Y w.r.t X 的雅可比矩阵并将X 的维度相加。然而，这在记忆方面是不可行的，因为我使用的函数是具有大量输入和输出的神经网络。

单独计算每个梯度（就像我在 cmets 中所做的那样）也是非常不可取的，因为这太慢了。

【问题讨论】：

你打算如何使用这样的结果张量？其背后的数学推理是什么？
我正在编写一个需要这种梯度掩码的 LRP（层相关性传播）实现。反正我觉得这个操作也不算太牵强。

标签： python pytorch autograd differentiation

【解决方案1】：

如果您每晚运行 pytorch，https://github.com/pytorch/pytorch/issues/10223 已部分实现，并且应该为大多数简单图表执行您想要的操作。您也可以尝试使用https://j-towns.github.io/2017/06/12/A-new-trick.html 中描述的技巧。

编辑：看起来https://pytorch.org/docs/stable/generated/torch.autograd.functional.jvp.html#torch.autograd.functional.jvp 为您实现了倒退技巧。所以你可以这样做：

from torch.autograd.functional import jvp
X = torch.eye(2)                              
X.requires_grad_()                            
def build_Y(x):                               
    return torch.sum(x*x, dim=0)              
                                              
print(jvp(build_Y, X, torch.ones(X.shape))[1])

【讨论】：