Skip to content

Instantly share code, notes, and snippets.

@qsh-zh
Forked from sbarratt/torch_jacobian.py
Created May 2, 2020 03:48
Show Gist options
  • Save qsh-zh/991950e1a7ed7562138408086c67f855 to your computer and use it in GitHub Desktop.
Save qsh-zh/991950e1a7ed7562138408086c67f855 to your computer and use it in GitHub Desktop.
Get the jacobian of a vector-valued function that takes batch inputs, in pytorch.
def get_jacobian(net, x, noutputs):
x = x.squeeze()
n = x.size()[0]
x = x.repeat(noutputs, 1)
x.requires_grad_(True)
y = net(x)
y.backward(torch.eye(noutputs))
return x.grad.data
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment