通过计算雅可比行列有效地使用 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只知道如何从标量张量传播梯度（我的网络的 output 不是），这就是我必须计算雅可比行列式的原因。

但是，这不是很有效并且有点慢，因为我的矩阵很大并且计算整个雅可比需要一段时间（而且我也没有使用整个雅可比矩阵）。

有没有办法只计算雅可比的对角线（在这个例子中得到 4 个向量）？

似乎有一个开放的功能请求，但似乎并没有引起太多关注。

更新1：
我尝试了@iacob 所说的关于设置torch.autograd.functional.jacobian(vectorize=True)的内容。
但是，这似乎更慢。 为了测试这一点，我将网络 output 从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-matrix 是非常低效的。

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

一点线性代数：

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可能是一长串矩阵-矩阵乘法的结果），那就太慢了。 相反，如果我们计算一个向量-雅可比乘法序列，计算会非常快。

---

解决方案

来自雅可比的样本元素：

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. 但如果我再次运行它，答案将会改变。

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

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

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

Answer 2

您是否尝试过设置torch.autograd.functional.jacobian(vectorize=True) ？

vectorize ( bool, optional ) - 此功能是实验性的，请自担风险使用。 在计算雅可比时，通常我们在雅可比的每一行调用一次autograd.grad 。 如果这个标志是True ，我们使用vmap原型特性作为后端来向量化对autograd.grad的调用，所以我们只调用它一次而不是每行一次。 这应该会在许多用例中带来性能改进，但是，由于此功能不完整，可能会出现性能悬崖。 请使用torch._C._debug_only_display_vmap_fallback_warnings(True)显示任何性能警告，如果您的用例存在警告，请向我们提交问题。 默认为False 。