Pytorch 对输入而不是输出求和 jacobian

Question

假设我有一个张量 Y（直接或间接）从张量 X.

计算得出

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

是否可以得到一个和Y一样形状的渐变mask来代替，其中每个元素mask[i][j]是[=20=的梯度之和] w.r.t。 X?

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

>>> 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 的维度求和即可。然而，这在内存方面是不可行的，因为我使用的函数是具有大量输入和输出的神经网络。

分别计算每个梯度（正如我在评论中所做的那样）也非常不可取，因为这太慢了。

Answer 1

如果你每晚运行 pytorch，https://github.com/pytorch/pytorch/issues/10223 is partially implemented and should do what you want for most simple graphs. You could also try using the trick described at 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])

Pytorch 对输入而不是输出求和 jacobian

Pytorch sum jacobian over inputs instead of outputs

python

differentiation

pytorch

autograd