Layer collection

Layer collections describe the structure of parameters that will be differentiated. We need the LayerCollection object in order to be able to map components of different objects together. As an example, when performing a matrix-vector product using a block diagonal representation, we need to make sure that elements of the vector corresponding to parameters from layer 1 are multiplied with the diagonal block also corresponding to parameters from layer 1, and so on.

Typical use cases include:

• All parameters of your network: for this you can simply use the constructor nngeometry.layercollection.LayerCollection.from_model()

• Only parameters from some layers of your network. In this case you need to

1. instantiate a new LayerCollection object

2. add your layers one at a time using nngeometry.layercollection.LayerCollection.add_layer_from_model()

class nngeometry.layercollection.AbstractLayer

class nngeometry.layercollection.Affine1dLayer(num_features, bias=True)

numel()

class nngeometry.layercollection.BatchNorm1dLayer(num_features)

numel()

class nngeometry.layercollection.BatchNorm2dLayer(num_features)

numel()

class nngeometry.layercollection.Conv2dLayer(in_channels, out_channels, kernel_size, bias=True)

numel()

class nngeometry.layercollection.ConvTranspose2dLayer(in_channels, out_channels, kernel_size, bias=True)

numel()

class nngeometry.layercollection.Cosine1dLayer(in_features, out_features)

numel()

class nngeometry.layercollection.GroupNormLayer(num_groups, num_channels)

numel()

class nngeometry.layercollection.LayerCollection(layers=None)

This class describes a set or subset of layers, that can be used in order to instantiate nngeometry.object.PVector or nngeometry.object.PSpaceDense objects

Parameters

layers –

add_layer(name, layer)

add_layer_from_model(model, module)

Add a layer by specifying the module corresponding to this layer (e.g. torch.nn.Linear or torch.nn.BatchNorm1d)

Parameters

• model – The model defining the neural network

• module – The layer to be added

from_model(ignore_unsupported_layers=False)

Constructs a new LayerCollection object by using all parameters of the model passed as argument.

Parameters

• model (nn.Module) – The PyTorch model

• ignore_unsupported_layers – If false, will raise an error

when model contains layers that are not supported yet. If true, will silently ignore the layer :type ignore_unsupported_layers: bool

get_layerid_module_maps(model)

numel()

Total number of scalar parameters in this LayerCollection object

Returns

number of scalar parameters

Return type

int

parameters(layerid_to_module)

class nngeometry.layercollection.LinearLayer(in_features, out_features, bias=True)

numel()

class nngeometry.layercollection.Parameter(*size)

numel()

class nngeometry.layercollection.WeightNorm1dLayer(in_features, out_features)

numel()

class nngeometry.layercollection.WeightNorm2dLayer(in_channels, out_channels, kernel_size)

numel()