[Python] Fast calculation of hessian of a neural network for part of the dataset

I need to calculate this formula.

This is the fastest way I found to calculate it

for i in range(number_parameters):
aux = self.loss_grad[:, i].detach()
product = torch.sum(aux * self.loss_grad[:, i])
first_term = torch.autograd.grad(product, self.model.parameters(), create_graph=True)
first_term = torch.cat([g.flatten() for g in first_term])
self.grad_sigma_diag = (1/(m-1)) * first_term


The idea is to make the scalar product between the elements of the datasets and then differenziate But since I have to do a for loop to calculate each entry, it is still slow.

How to solve this?

Continue reading...