python - 通过计算雅可比行列式有效地使用 PyTorch 的 autograd 与张量

Question

在我之前的 question我找到了如何将 PyTorch 的 autograd 与张量一起使用:

import torch
from torch.autograd import grad
import torch.nn as nn
import torch.optim as optim

class net_x(nn.Module): 
        def __init__(self):
            super(net_x, self).__init__()
            self.fc1=nn.Linear(1, 20) 
            self.fc2=nn.Linear(20, 20)
            self.out=nn.Linear(20, 4) #a,b,c,d

        def forward(self, x):
            x=torch.tanh(self.fc1(x))
            x=torch.tanh(self.fc2(x))
            x=self.out(x)
            return x

nx = net_x()

#input
t = torch.tensor([1.0, 2.0, 3.2], requires_grad = True) #input vector
t = torch.reshape(t, (3,1)) #reshape for batch

#method 
dx = torch.autograd.functional.jacobian(lambda t_: nx(t_), t)
dx = torch.diagonal(torch.diagonal(dx, 0, -1), 0)[0] #first vector
#dx = torch.diagonal(torch.diagonal(dx, 1, -1), 0)[0] #2nd vector
#dx = torch.diagonal(torch.diagonal(dx, 2, -1), 0)[0] #3rd vector
#dx = torch.diagonal(torch.diagonal(dx, 3, -1), 0)[0] #4th vector
dx 
>>> 
tensor([-0.0142, -0.0517, -0.0634])

问题是 grad只知道如何从标量张量(我的网络的输出不是)传播梯度，这就是我必须计算雅可比行列式的原因。
但是，这不是很有效并且有点慢，因为我的矩阵很大并且计算整个雅可比矩阵需要一段时间(而且我也没有使用整个雅可比矩阵)。
有没有办法只计算雅可比矩阵的对角线(在这个例子中得到 4 个向量)？
似乎有一个 open feature request但它似乎并没有受到太多关注。
更新 1:
我尝试了@iacob 所说的设置 torch.autograd.functional.jacobian(vectorize=True) .
然而，这似乎更慢。为了测试这个，我改变了我的网络输出 4至 400 ，和我的输入 t成为:

val = 100
t = torch.rand(val, requires_grad = True) #input vector
t = torch.reshape(t, (val,1)) #reshape for batch

无 vectorized = True :

Wall time: 10.4 s

和:

Wall time: 14.6 s

Answer 1

好，先出结果:
性能(我的笔记本电脑有一个 RTX-2070，PyTorch 正在使用它):

# Method 1: Use the jacobian function
CPU times: user 34.6 s, sys: 165 ms, total: 34.7 s
Wall time: 5.8 s

# Method 2: Sample with appropriate vectors
CPU times: user 1.11 ms, sys: 0 ns, total: 1.11 ms
Wall time: 191 µs

它快了大约 30000 倍。

为什么要使用 backward而不是 jacobian (在你的情况下)
我不是 PyTorch 的专家。但是，根据我的经验，如果不需要其中的所有元素，计算 jacobi 矩阵的效率非常低。
如果你只需要对角线元素，你可以使用 backward计算函数 矢量 -jacobian 乘法与一些特定的向量。如果您设置了 矢量 s 正确，您可以从雅可比矩阵中采样/提取特定元素。
一点线性代数:

j = np.array([[1,2],[3,4]]) # 2x2 jacobi you want 
sv = np.array([[1],[0]])     # 2x1 sampling vector

first_diagonal_element = sv.T.dot(j).dot(sv)  # it's j[0, 0]

对于这个简单的案例，它并不是那么强大。但是，如果 PyTorch 需要沿途计算所有雅可比( j 可能是一长串矩阵-矩阵乘法的结果)，那就太慢了。相比之下，如果我们计算一系列向量雅可比乘法，则计算速度会非常快。

解决方案
来自 jacobian 的示例元素:

import torch
from torch.autograd import grad
import torch.nn as nn
import torch.optim as optim

class net_x(nn.Module): 
        def __init__(self):
            super(net_x, self).__init__()
            self.fc1=nn.Linear(1, 20) 
            self.fc2=nn.Linear(20, 20)
            self.out=nn.Linear(20, 400) #a,b,c,d

        def forward(self, x):
            x=torch.tanh(self.fc1(x))
            x=torch.tanh(self.fc2(x))
            x=self.out(x)
            return x

nx = net_x()

#input

val = 100
a = torch.rand(val, requires_grad = True) #input vector
t = torch.reshape(a, (val,1)) #reshape for batch


#method 
%time dx = torch.autograd.functional.jacobian(lambda t_: nx(t_), t)
dx = torch.diagonal(torch.diagonal(dx, 0, -1), 0)[0] #first vector
#dx = torch.diagonal(torch.diagonal(dx, 1, -1), 0)[0] #2nd vector
#dx = torch.diagonal(torch.diagonal(dx, 2, -1), 0)[0] #3rd vector
#dx = torch.diagonal(torch.diagonal(dx, 3, -1), 0)[0] #4th vector
print(dx)


out = nx(t)
m = torch.zeros((val,400))
m[:, 0] = 1
%time out.backward(m)
print(a.grad)

a.grad应该等于第一个张量 dx .而且， m是与代码中所谓的“第一个向量”相对应的采样向量。

1. but if I run it again the answer will change.

是的，你已经想通了。每次调用 backward 都会累积梯度.所以你必须设置 a.grad如果您必须多次运行该单元格，请先归零。

2. can you explain the idea behind the m method? Both using the torch.zeros and setting the column to 1. Also, how come the grad is on a rather than on t?

m背后的理念方法是:什么函数backward计算实际上是一个向量雅可比乘法，其中向量代表所谓的“上游梯度”，而雅可比矩阵是“局部梯度”(这个雅可比也是你用jacobian函数得到的，因为您的 lambda 可以被视为单个“本地”操作)。如果您需要 jacobian 中的一些元素，您可以伪造(或更准确地说，构造)一些“上游梯度”，您可以使用它从 jacobian 中提取特定条目。但是，如果涉及复杂的张量运算，有时可能很难找到这些上游梯度(至少对我而言)。

PyTorch 在计算图的叶节点上累积梯度。而且，您的原始代码行 t = torch.reshape(t, (3,1))失去叶节点的句柄，并且 t现在指的是中间节点而不是叶节点。为了访问叶节点，我创建了张量 a .