[ADD] Custom slicing module + support for some 2nd-order extensions

f-dangel / backpack

BackPACK - a backpropagation package built on top of PyTorch which efficiently computes quantities other than the gradient.

https://backpack.pt/

MIT License

555 stars 55 forks source link

[ADD] Custom slicing module + support for some 2nd-order extensions #241

Closed f-dangel closed 2 years ago

f-dangel commented 2 years ago

There is no nn.Module in torch to perform tensor slicing, e.g. to implement certain ResNets for CIFAR10. This PR adds a custom slicing module with support for some second-order extensions.

Add Slicing module
Implement jac_t_mat_prod
Support {Batch}DiagGGN{Exact, MC}, {Batch}DiagHessian, SqrtGGN{Exact, MC} extensions