Hi, I have a simple question: does anyone know a straightforward method to compute both the gradient and hessian of a scalar function simultaneously using autograd? For example, when you use the torch.hessian function, the gradient is computed during this process - how do I get access to it?
In the standalone autograd library (outside of pytorch) I had to edit the sourcecode in order to do this, I was hoping I wouldn’t have to do the same here!
Thank you to anyone who replies, Sean.
Hi @FOXP20,
You can use torch.func package to get this. You can simply define a function, which computes the derivative and create a copy of it. Then you can compute the derivative again (and return the copy of the 1st derivative at the same time as an auxiliary variable).
Here’s a minimal reproducible example,
import torch
from torch.func import jacrev #reverse-mode AD
def func(x):
return x**2
def jacobian_func(x):
jac = jacrev(func, argnums=(0))(x)
return jac, jac #return 2 copies here
def hessian_with_grad_func(x):
hess, jac = jacrev(jacobian_func, argnums=(0), has_aux=True)(x)
return hess, jac
x = torch.randn(1) #if you have more samples, you'll need to use torch.func.vmap
hess, jac = hessian_with_grad_func(x)
print("output: ",x)
print("jacobian: ",jac) #equals 2*x
print("hessian: ",hess) #equals 2
Fantastic that was exactly the type of response I was looking for - thank you!