Gradients in loss function and efficient calculation

1

I’m using the l2 norm of a pde as the loss function, the input are variables (x,y) and the output are some quantities (r, u, v), and I want to calculate the jacobian dr/dx, dr/dy, du/dx … and I’d like to do it in batch - (I believe it’s very slow otherwise, though cant be sure since I haven’t been able to try!). Currently I use the jacobian function and loop over each sample in the batch, as shown below. I’d like to do this without the for loop (or much more efficiently) can this be done as a simple batch operation? If so how?

    x, y = batch
    x.requires_grad_(True)
    y_hat = self(x)

    jacobian_list = []
    
   for inp in x:
        inp = inp.unsqueeze(dim=0)
        jacobian = torch.autograd.functional.jacobian(
            self, inp, create_graph=True, strict=False, vectorize=False
        )
        jacobian_list.append(jacobian)

    gradients = torch.reshape(torch.stack(jacobian_list), (-1, 3, 2))

Thanks!

2

Looks like this has also been discussed here Computing batch Jacobian efficiently - #4 by albanD, not sure if there are any updates.

3

If you want batched Jacobian (or higher order) gradient functions you can check out functorch, their github page is here.

1 Like