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

has_bias()

transposed = False

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.Conv1dLayer(in_channels, out_channels, kernel_size, bias=True)

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.EmbeddingLayer(num_embeddings, embedding_dim)

numel()

transposed = True

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_common_layers(other)

get_layerid_module_map(model)

merge(other)

named_parameters(layerid_to_module)

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.LayerNormLayer(normalized_shape, bias=True)

numel()

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

numel()

class nngeometry.layercollection.Parameter(*size)

numel()

class nngeometry.layercollection.RMSNormLayer(normalized_shape)

numel()

class nngeometry.layercollection.WeightNorm1dLayer(in_features, out_features)

numel()

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

numel()