pytorch神经网络的前向Jacobian速度慢

时间:2019-01-26 22:40:13

标签: python neural-network pytorch

我正在计算pytorch中2层前馈神经网络的正向雅可比(输出相对于输入的导数),我的结果是正确的,但相对较慢。考虑到计算的性质,我希望它的速度大约与通过网络的正向传递速度一样快(或者可能是网络的2-3倍),但是在此例程上运行优化步骤则需要大约12倍的时间(在我的测试示例我只想在所有点上都使jacobian = 1)与标准均方误差,所以我认为我正在以一种非最佳方式进行操作。我只是想知道是否有人知道一种更快的编码方式。我的测试网络有2个输入节点,其后是2个隐藏层(每个节点5个)和2个节点的输出层,并在隐藏层上使用tanh激活函数,并具有线性输出层。

Jacobian计算基于论文The Limitations of Deep Learning in Adversarial Settings,该论文给出了前向导数的基本递归定义(基本上,您最终将激活函数的导数与各层的权重和先前的偏导数相乘)。这与正向传播非常相似,这就是为什么我希望它比现在快。然后,最后2x2 jacobian的行列式非常简单。

下面是网络和雅各布的代码

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

,这是被比较的两个错误函数。请注意,第一个函数需要通过网络进行额外的前向调用,以获取输出值(标记为动作),而第二个函数则不需要,因为它可以处理输入值。

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

任何对此的见识将不胜感激

1 个答案:

答案 0 :(得分:0)

第二个“ a”计算在我的计算机(cpu)上花费的时间最多。

# Here you increase the size of the matrix with a factor of "input_1"
expanded_deriv = tanh_deriv_tensor.unsqueeze(-1).expand(-1, -1, input_1)
partials = expanded_deriv * a.expand_as(expanded_deriv)

# Here your torch.matmul() needs to handle "input_1" times more computations than in a normal forward call
a = torch.matmul(self.h_1_2.weight, partials)

在我的机器上,计算雅可比行列式的时间大约是用火炬计算的时间

a = torch.rand(hidden_1, hidden_2)
b = torch.rand(n_inputs, hidden_1, input_1)
%timeit torch.matmul(a,b)

在计算方面,我认为不可能加快速度。除非您可以从CPU转移到GPU,否则在大型矩阵上GPU会变得更好。