Forward Jacobian Of Neural Network in Pytorch is Slow
I am computing the forward jacobian (derivative of outputs with respect to inputs) of a 2 layer feedforward neural network in pytorch, and my results are correct but relatively slow. Given the nature of the calculation I would expect it to be approximately as fast as a forward pass through the network (or maybe 2-3x as long), but it takes ~12x as long to run an optimization step on this routine (in my test example I just want the jacobian=1 at all points) vs the standard mean squared error so I assume I am doing something in an un-optimal manner. I'm just wondering if anyone knew a faster way to code this. My test network has 2 input nodes, followed by 2 hidden layers of 5 nodes each and an output layer of 2 nodes, and uses tanh activation functions on the hidden layers, with a linear output layer. 

The Jacobian calculations are based on the paper The Limitations of Deep Learning in Adversarial Settings which gives a basic recursive definition of the forward derivative (basically you end up multiplying the derivative of your activation functions with the weights and previous partial derivatives of each layer). This is very similar to forward propagation, which is why I would expect it to be faster than it is. Then the determinant of the 2x2 jacobian at the end is pretty straightforward. 

Below is the code for the network and the 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


and here are the two error functions being compared. Note that the first one requires an extra forward call through the network, to get the output values (labeled action) while the second function does not since it works on the input values.

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


Any insight on this would be greatly appreciated

The second calculation of 'a' takes the most time on my machine (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)


On my machine the time of computing the Jacobian is roughly the time it takes torch to compute

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


I don't think it is possible to speed this up, computationally wise. Unless you can move from CPU to GPU, because GPU get better on larges matrices. 
