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

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 个向量）？

似乎有一个开放的功能请求但似乎没有得到太多关注。

更新 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

好的，先看结果：

\n

性能（我的笔记本电脑有 RTX-2070，PyTorch 正在使用它）：

\n

# Method 1: Use the jacobian function\nCPU times: user 34.6 s, sys: 165 ms, total: 34.7 s\nWall time: 5.8 s\n\n# Method 2: Sample with appropriate vectors\nCPU times: user 1.11 ms, sys: 0 ns, total: 1.11 ms\nWall time: 191 \xc2\xb5s\n

\n

速度快了大约 30000 倍。

\n\n

为什么你应该使用backward而不是jacobian（在你的情况下）

\n

我不是 PyTorch 的专业人士。但是，根据我的经验，如果不需要雅可比矩阵中的所有元素，那么计算雅可比矩阵的效率相当低。

\n

如果只需要对角线元素，可以使用backward函数来计算向量雅可比与某些特定向量的乘法。如果正确设置向量s，则可以从雅可比矩阵中采样/提取特定元素。

\n

一点线性代数：

\n

j = np.array([[1,2],[3,4]]) # 2x2 jacobi you want \nsv = np.array([[1],[0]])     # 2x1 sampling vector\n\nfirst_diagonal_element = sv.T.dot(j).dot(sv)  # it\'s j[0, 0]\n

\n

对于这个简单的例子来说，它的功能并不是那么强大。但如果 PyTorch 需要计算一路上的所有雅可比矩阵（j可能是一长串矩阵-矩阵乘法的结果），那就太慢了。相反，如果我们计算向量雅可比乘法序列，计算速度将会非常快。

\n\n

解决方案

\n

雅可比的示例元素：

\n

import torch\nfrom torch.autograd import grad\nimport torch.nn as nn\nimport torch.optim as optim\n\nclass net_x(nn.Module): \n        def __init__(self):\n            super(net_x, self).__init__()\n            self.fc1=nn.Linear(1, 20) \n            self.fc2=nn.Linear(20, 20)\n            self.out=nn.Linear(20, 400) #a,b,c,d\n\n        def forward(self, x):\n            x=torch.tanh(self.fc1(x))\n            x=torch.tanh(self.fc2(x))\n            x=self.out(x)\n            return x\n\nnx = net_x()\n\n#input\n\nval = 100\na = torch.rand(val, requires_grad = True) #input vector\nt = torch.reshape(a, (val,1)) #reshape for batch\n\n\n#method \n%time dx = torch.autograd.functional.jacobian(lambda t_: nx(t_), t)\ndx = torch.diagonal(torch.diagonal(dx, 0, -1), 0)[0] #first vector\n#dx = torch.diagonal(torch.diagonal(dx, 1, -1), 0)[0] #2nd vector\n#dx = torch.diagonal(torch.diagonal(dx, 2, -1), 0)[0] #3rd vector\n#dx = torch.diagonal(torch.diagonal(dx, 3, -1), 0)[0] #4th vector\nprint(dx)\n\n\nout = nx(t)\nm = torch.zeros((val,400))\nm[:, 0] = 1\n%time out.backward(m)\nprint(a.grad)\n

\n

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

\n\n

\n

\n
1. 但如果我再次运行它，答案就会改变。
\n

\n

\n

是的，你已经明白了。每次调用时梯度都会累积backward。a.grad因此，如果您必须多次运行该单元，则必须首先设置为零。

\n

\n

\n
2. 你能解释一下该m方法背后的想法吗？都使用torch.zeros并将列设置为1。另外，为什么 grad 是 ona而不是 on t？
\n

\n

\n

\n
• 该方法背后的思想m是：函数backward计算的实际上是向量雅可比矩阵乘法，其中向量代表所谓的“上游梯度”，雅可比矩阵就是“局部梯度”（而这个雅可比矩阵也是您可以使用该jacobian功能，因为您lambda可以被视为单个“本地”操作）。如果您需要来自雅可比的一些元素，您可以伪造（或者更准确地说，构造）一些“上游梯度”，用它您可以从雅可比中提取特定的条目。然而，有时如果涉及复杂的张量运算，这些上游梯度可能很难找到（至少对我来说）。
\n
• PyTorch 在计算图的叶节点上累积梯度。而且，您原来的代码t = torch.reshape(t, (3,1))行失去了叶节点的句柄，t现在引用的是中间节点而不是叶节点。为了访问叶节点，我创建了张量a。
\n

\n