Python Pytorch中神经网络的前向雅可比矩阵速度较慢
我正在Pytork中计算一个2层前馈神经网络的前向雅可比矩阵(输出对输入的导数),我的结果是正确的,但速度相对较慢。考虑到计算的性质,我希望它大约与通过网络的前向传递一样快(或者可能是2-3倍长),但在这个例程上运行优化步骤需要12倍的时间(在我的测试示例中,我只希望雅可比数在所有点上都为1)vs标准均方误差,所以我假设我在以非最佳方式做某事。我只是想知道是否有人知道一种更快的编码方法。我的测试网络有2个输入节点,然后是2个隐藏层,每个隐藏层有5个节点,输出层有2个节点,并在隐藏层上使用tanh激活函数,带有线性输出层 雅可比计算是基于一篇论文,该论文给出了正向导数的基本递归定义(基本上,你最终将激活函数的导数乘以每层的权重和之前的偏导数)。这与前向传播非常相似,这就是为什么我希望它比现在快的原因。最后2x2雅可比矩阵的行列式非常简单 下面是网络和雅可比矩阵的代码Python Pytorch中神经网络的前向雅可比矩阵速度较慢,python,neural-network,pytorch,Python,Neural Network,Pytorch,我正在Pytork中计算一个2层前馈神经网络的前向雅可比矩阵(输出对输入的导数),我的结果是正确的,但速度相对较慢。考虑到计算的性质,我希望它大约与通过网络的前向传递一样快(或者可能是2-3倍长),但在这个例程上运行优化步骤需要12倍的时间(在我的测试示例中,我只希望雅可比数在所有点上都为1)vs标准均方误差,所以我假设我在以非最佳方式做某事。我只是想知道是否有人知道一种更快的编码方法。我的测试网络有2个输入节点,然后是2个隐藏层,每个隐藏层有5个节点,输出层有2个节点,并在隐藏层上使用tanh
class Network(torch.nn.Module):
def __init__(self):
super(Network, self).__init__()
self.h_1_1 = torch.nn.Linear(input_1, hidden_1)
self.h_1_2 = torch.nn.Linear(hidden_1, hidden_2)
self.out = torch.nn.Linear(hidden_2, out_1)
def forward(self, x):
x = F.tanh(self.h_1_1(x))
x = F.tanh(self.h_1_2(x))
x = (self.out(x))
return x
def jacobian(self, x):
a = self.h_1_1.weight
x = F.tanh(self.h_1_1(x))
tanh_deriv_tensor = 1 - (x ** 2)
expanded_deriv = tanh_deriv_tensor.unsqueeze(-1).expand(-1, -1, input_1)
partials = expanded_deriv * a.expand_as(expanded_deriv)
a = torch.matmul(self.h_1_2.weight, partials)
x = F.tanh(self.h_1_2(x))
tanh_deriv_tensor = 1 - (x ** 2)
expanded_deriv = tanh_deriv_tensor.unsqueeze(-1).expand(-1, -1, out_1)
partials = expanded_deriv*a
partials = torch.matmul(self.out.weight, partials)
determinant = partials[:, 0, 0] * partials[:, 1, 1] - partials[:, 0, 1] * partials[:, 1, 0]
return determinant
下面是比较的两个误差函数。请注意,第一个函数需要通过网络进行额外的前向调用,以获取输出值(标记为action),而第二个函数则不需要,因为它处理输入值
def actor_loss_fcn1(action, target):
loss = ((action-target)**2).mean()
return loss
def actor_loss_fcn2(input): # 12x slower
jacob = model.jacobian(input)
loss = ((jacob-1)**2).mean()
return loss
如果对此有任何见解,我们将不胜感激。第二次计算“a”在我的机器(cpu)上花费的时间最多 在我的机器上,计算雅可比矩阵的时间大约是计算torch所需的时间
a = torch.rand(hidden_1, hidden_2)
b = torch.rand(n_inputs, hidden_1, input_1)
%timeit torch.matmul(a,b)
从计算的角度来看,我认为不可能加快速度。除非你能从CPU转移到GPU,因为GPU在大型矩阵上表现得更好
a = torch.rand(hidden_1, hidden_2)
b = torch.rand(n_inputs, hidden_1, input_1)
%timeit torch.matmul(a,b)