Pythorch中神经网络的前向雅可比矩阵

Question

我正在Pythorch中计算一个2层前向神经网络的前向雅可比（输出相对于输入的导数），我的结果是正确的，但相对较慢。考虑到计算的性质，我预计它的速度大约与通过网络的前向传输一样快（或者可能是2-3倍长），但是在这个程序上运行优化步骤所花的时间是标准均方误差的12倍（在我的测试示例中，我只希望jacobian=1在所有点上都是1），所以我假设我在以一种非最优的方式做一些事情。我只是想知道有没有人知道更快的编码方法。我的测试网络有2个输入节点，接着是2个隐藏层，每个隐藏层有5个节点，输出层有2个节点，在隐藏层上使用tanh激活函数，输出层是线性的。在

雅可比计算是基于论文The Limitations of Deep Learning in Adversarial Settings给出了一个基本的前向导数的递归定义（基本上，你最终将激活函数的导数乘以每个层的权重和以前的偏导数）。这与前向传播非常相似，这就是为什么我希望它比现在更快。最后的2x2雅可比行列式就很简单了。在

下面是网络和雅可比矩阵的代码

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），而第二个函数则不需要，因为它处理输入值。在

^{pr2}$

对此有任何见解将不胜感激

Answer 1

第二次计算'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)

在我的机器上，计算雅可比的时间大概就是火炬计算所需的时间

^{pr2}$

我认为不可能在计算上加快速度。除非你能从一个CPU移到另一个GPU，因为GPU在大矩阵上表现得更好。在