[ADD] Custom padding module + support 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 padding module + support some 2nd-order extensions #243

Closed f-dangel closed 2 years ago

f-dangel commented 2 years ago

Adds a custom Pad module which is the nn.Module equivalent of torch.nn.functional.pad. This can be useful for implementing certain ResNets for CIFAR10. This PR also adds support for some second-order extensions.

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