Skip to content

Instantly share code, notes, and snippets.

@apaszke
Last active April 3, 2024 03:40
Show Gist options
  • Save apaszke/226abdf867c4e9d6698bd198f3b45fb7 to your computer and use it in GitHub Desktop.
Save apaszke/226abdf867c4e9d6698bd198f3b45fb7 to your computer and use it in GitHub Desktop.
import torch
def jacobian(y, x, create_graph=False):
jac = []
flat_y = y.reshape(-1)
grad_y = torch.zeros_like(flat_y)
for i in range(len(flat_y)):
grad_y[i] = 1.
grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
jac.append(grad_x.reshape(x.shape))
grad_y[i] = 0.
return torch.stack(jac).reshape(y.shape + x.shape)
def hessian(y, x):
return jacobian(jacobian(y, x, create_graph=True), x)
def f(x):
return x * x * torch.arange(4, dtype=torch.float)
x = torch.ones(4, requires_grad=True)
print(jacobian(f(x), x))
print(hessian(f(x), x))
@guanshaoheng
Copy link

guanshaoheng commented Apr 17, 2021

Now the function torch.autograd.functional.jacobian can do the same thing, I think.


def jacobian(y, x, create_graph=False):
    # xx, yy = x.detach().numpy(), y.detach().numpy()
    jac = []
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    for i in range(len(flat_y)):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=True)
        jac.append(grad_x.reshape(x.shape))
        grad_y[i] = 0.
    return torch.stack(jac).reshape(y.shape + x.shape)


def hessian(y, x):
    return jacobian(jacobian(y, x, create_graph=True), x)


def f(xx):
    # y = x * x * torch.arange(4, dtype=torch.float)
    matrix = torch.tensor([[0.2618, 0.2033, 0.7280, 0.8618],
        [0.1299, 0.6498, 0.6675, 0.0527],
        [0.3006, 0.9691, 0.0824, 0.8513],
        [0.7914, 0.2796, 0.3717, 0.9483]], requires_grad=True)
    y = torch.einsum('ji, i -> j', (matrix, xx))
    return y


if __name__ == "__main__":
    # matrix = torch.rand(4, 4, requires_grad=True)
    # print(matrix)
    x = torch.arange(4,  dtype=torch.float, requires_grad=True)
    print(jacobian(f(x), x))
    grad = torch.autograd.functional.jacobian(f, x).numpy()
    # grad = grad.flatten()
    print(grad)
    # print(hessian(f(x, matrix), x))

output

        [0.1299, 0.6498, 0.6675, 0.0527],
        [0.3006, 0.9691, 0.0824, 0.8513],
        [0.7914, 0.2796, 0.3717, 0.9483]], grad_fn=<ViewBackward>)
[[0.2618 0.2033 0.728  0.8618]
 [0.1299 0.6498 0.6675 0.0527]
 [0.3006 0.9691 0.0824 0.8513]
 [0.7914 0.2796 0.3717 0.9483]]```

@AjinkyaBankar
Copy link

Hi,
I want to find a Hessian matrix for the loss function of the pre-trained neural network with respect to the parameters of the network. How can I use this method? Can someone please share an example? Thanks.

@maryamaliakbari
Copy link

Hi,
I want to find a Hessian matrix for the loss function of the pre-trained neural network with respect to the parameters of the network. How can I use this method? Can someone please share an example? Thanks.

Hi,
I am looking for the same thing. Could you figure out how we can do it?

@mil-ad
Copy link

mil-ad commented Oct 25, 2021

I think this has now been added to recent versions of torch's autograd module. Maybe look at the examples here

@maryamaliakbari
Copy link

I think this has now been added to recent versions of torch's autograd module. Maybe look at the examples here

Right. I checked it. When I use this method I am getting multiple errors. I am looking for an example or similar code to see how the implementation is done.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment