Efficient way of Calculating jacobians

Hello everyone.
I would like to take the derivative of the output with respect to the input. People in other pages have suggested this:

torch.autograd.grad(output, input, retain_graph=True)[0]

However, I think this method is not very efficient: Many of the operations needed to take derivative wrt inputs have already been calculated during the training (going backward from output to the layer just before the input). So, if, after training, I run the code above, I’m basically repeating many operations. Is that right? Is there any way that I take the derivatives wrt inputs using the last update of the training?

Thank you very much in advance.

Hi,

You can make your input require gradients. That way the backward() call will populate its .grad field with it’s gradients !

Thanks a lot.

But doesn’t the model update the inputs in each iteration by making the inputs require grads? Because the model cannot distinguish between parameters and inputs in this way, right?

Also, would you please write me the code that extracts these gradients w.r.t the inputs from backward()?
Thank you very much.

The .backward() computes the gradients, it’s the optimizer .step() that do update of the weights that were given to it when it is created.

To get the gradients you can do:

inp = # Some input
label = # Some label
inp.requires_grad_()

out = model(inp)
loss = criterion(out, label)

# Do a regular gradient step from the optimizer
optimizer.zero_grad()
loss.backward()
optimizer.step()

# Grab input gradient
inp_grad = inp.grad

Thank you very much.
I ran this code, but it’s not giving me the right answer. To test this code, I ran a linear regression: y= w0+w1*x, and then I grabed grad of y w.r.t. x using x.grad.
It should give me an n-long tensor with estimate of w1 on each entry of the tensor. Right? But it gives me this:

tensor([[ 0.1268],
[ 0.2045],
[-0.3171],
…,
[ 0.3761],
[ 0.1287],
[-0.1889]], device=‘cuda:0’)

Do you know what’s going on? Maybe x.grad doesn’t give derivative of y w.r.t. x?

Hi,

It works as expected:

import torch 

inp = torch.rand(5, 10, requires_grad=True)

w1 = torch.rand(10, 1, requires_grad=True)
w0 = torch.rand(1, 1, requires_grad=True)

out = w0 + inp.mm(w1)

out.sum().backward()

print("w1")
print(w1)

print("inp.grad")
print(inp.grad)

Keep in mind here that the first dimension of inp is a “batch” and so the gradient will be w1 repeated as many times as they are elements in the batch.

1 Like

Hi,

Thank you. You code gives the right answer, but if you do it my way, it doesn’t give the correct answer. Please run this code to see what I mean. Thanks a lot.

x = torch.rand(10, 1, requires_grad=False)#simulate x
y = 2 + 3 * x + .1 * torch.rand(10, 1)#simulate y from linear regression

Run linear regression

class net(torch.nn.Module):
    def __init__(self):
        super(net, self).__init__()

        self.linear = torch.nn.Linear(1, 1, bias=True)

    def forward(self, x):
        out = self.linear(x)
        return(out)


x.requires_grad_()
# training
model = net()

model.train()

loss = torch.nn.MSELoss(size_average=True)
optimizer  = torch.optim.SGD(model.parameters(), lr = 0.01)
 

for epoch in range(10):
    
    y_pred = model(x)

    l = loss(y_pred, y)

    optimizer.zero_grad()
    l.backward()
    optimizer.step()

print(x.grad)

The main difference I see is that you use MSE to compute the loss, that means that the gradients you will get are not just the weights of the linear but a function of y as well.

Hi.
That’s right. But how can I get dy/dx, in the training process? I mean, apply backward() on the loss function but somehow get dy/dx? Because going back thourgh the computation graph of a simple one layer NN, backpropagations calculates these:
dL/dy
dy/dH
dH/dW
. Right?
Since x.reauires_grad_() was applied, we also have dH/dx. So we should somehow get dy/dx=dy/dH * dH/dx, automatically, but I don’t know how to do this.

As I said before, I know the following code works, but I’m looking for something that doesn’t need to take all the gradients one more time after training is completed.

torch.autograd.grad(output, input, retain_graph=True)[0]

Thanks a lot.

If I understand your case, you want something like:

input = # Some input
output = net(input)
loss = crit(output, labels)

The think the that loss.backward will update the weights in the network based on the loss value. If you want gradients of just output wrt inputs, then you only want the influence of output wrt to input and not the other weights in the network.
This means that you cannot do both in one single backward pass as you want in one case to consider the influence on the weights (for the loss) and you don’t want it on the other case (for the grad op).

Another point is that this autograd,grad call won’t give you the full jacobian of the network but the vector jacobian product with a vector of all ones. Or you can see it as the full jacobian of the function that does net(input).sum().

Thank you very much :slight_smile: