search

tfjgeorge / nngeometry

{KFAC,EKFAC,Diagonal,Implicit} Fisher Matrices and finite width NTKs in PyTorch
https://nngeometry.readthedocs.io
MIT License
206 stars 20 forks source link

Update layercollection.py #15

Closed milpan closed 3 years ago

milpan commented 3 years ago

Updated to allow for R2 Convolutional layers

tfjgeorge commented 3 years ago

Hello, please do not hesitate to reach if you need help with implementing custom layers.

milpan commented 3 years ago

Hello, please do not hesitate to reach if you need help with implementing custom layers.

Hey there mate! Sorry i didn't see your message a month ago!!! Really love this library, had a play around with it on some convolutional networks and the results were great! I'm working on obtaining FIM for the E2CNN library, at the moment ive managed to pull the fisher info out for the last fully connected layer of the network. Was trying to get it working with the R2 Conv Layer, but i think there's a big difference in implementation.

Also, cant thank u enough for reaching out! Kind Regards Matt

tfjgeorge commented 3 years ago

Hey!

If you want to add a new layer type, you need to implement per-example gradient computation in the generator, given the input of the layer, and the gradient w.r.t the output of the layer.

For linear layers it is quite easy (see e.g. https://github.com/tfjgeorge/nngeometry/blob/master/nngeometry/generator/jacobian.py#L802-L806), but it gets more complicated for more sophisticated ones.

tfjgeorge commented 3 years ago

I am closing this PR. Please reopen it if you want to implement per-example gradient computation for your new layers.