search

cybertronai / autograd-hacks

The Unlicense
153 stars 32 forks source link

Can compute_grad1 support multiple backprops? #6

Open zhangxz1123 opened 3 years ago

zhangxz1123 commented 3 years ago

Hi Yaroslav, First of all, I want to say your tricks here are really neat and have helped me a lot!

I'm currently trying to compute the per-example gradient over two backprops in a typical Hessian vector product calculation, e.g.

output = model(data)
loss = loss_fn(output, targets)
grad = torch.autograd.grad(loss, model.parameters(), create_graph=True)
grad = flat_grad(grad)
grad_p = (grad * p).sum()
hessian_p= torch.autograd.grad(grad_p , actor.parameters())

Is it possible to apply compute_grad1 to this procedure to get the per-example Hessian vector product? Naively applying it fails because currently compute_grad1 only supports a single backprop.

Thanks! Jack