Getting the output's grad with respect to the input

Question

I'm currently trying to implement an ODE Solver with Pytorch, my solution requires computing the gradient of each output wtr to its input.

y = model(x)

for i in range(len(y)): #compute output grad wrt input
       y[i].backward(retain_graph=True)
    
ydx=x.grad 

I was wondering if there is a more elegant way to compute the gradients for each output in the batch, since the code gets messy for higher order ODEs and PDEs. I tried using:

torch.autograd.backward(x,y,retain_graph=True)

without much success.

Answer 1

Try the torch.autograd.functional.jacobian if your version of PyTorch has the API implemented. I am doing the same for the Burger's equation, and posted this thread on the same subject: PyTorch how to compute second order jacobian?

Solving PDE with DL is a hot topic right now

Answer 2

You can use torch.autograd.grad function to obtain gradients directly. One problem is that it requires the output (y) to be scalar. Since your output is an array, you will still need to loop through its values.

The call will look something like this.

[torch.autograd.grad(outputs=out, inputs=x, retain_graph=True)[0][i] 
    for i, out in enumerate(y)]

Here is an example of what I mean. Let's consider variable x with values [1,2,3] and a model that just squares its input.

x = torch.Tensor([1, 2, 3])
x.requires_grad = True

def model(x): return x ** 2

y = model(x)

Now, if you call the torch.autograd.grad as I have described, you will get:

[torch.autograd.grad(outputs=out, inputs=x, retain_graph=True)[0][i] 
    for i, out in enumerate(y)]

# [tensor(2.), tensor(4.), tensor(6.)]

Which is the list of derivatives wrt. to values of x - [ydx0, ydx1, ydx2]