Generators

The spirit of NNGeometry is that you do not directly manipulate Generator objects, these can be considered as a backend that you do not need to worry about once instantiated. You instead instantiate concrete representations such as PMatDense or PMatKFAC and directly call linear algebra operations on these concrete representations.

class nngeometry.generator.jacobian.Jacobian(model, function=None, n_output=1, centering=False, layer_collection=None)

Computes jacobians \(\mathbf{J}_{ijk}=\frac{\partial f\left(x_{j}\right)_{i}}{\delta\mathbf{w}_{k}}\), FIM matrices \(\mathbf{F}_{k,k'}=\frac{1}{n}\sum_{i,j}\mathbf{J}_{ijk}\mathbf{J}_{ijk'}\) and NTK matrices \(\mathbf{K}_{iji'j'}=\sum_{k}\mathbf{J}_{ijk}\mathbf{J}_{ijk'}\).

This generator is written in pure PyTorch and exploits some tricks in order to make computations more efficient.

Parameters
  • layer_collection (layercollection.LayerCollection) –

  • model (Pytorch nn.Module) –

  • function (python function) – A function \(f\left(X,Y,Z\right)\) where \(X,Y,Z\) are minibatchs returned by the dataloader (Note that in some cases \(Y,Z\) are not required). If None, it defaults to function = lambda *x: model(x[0])

  • n_output (integer) – How many output is there for each example of your function. E.g. in 10 class classification this would probably be 10.