models

A collection of equivariant neural network architectures.

class symm_learning.models.EMLP(in_type: FieldType, out_type: FieldType, hidden_units: int = 128, activation: str | list[str] = 'ReLU', pointwise_activation: bool = True, bias: bool = True, hidden_rep: Representation = None)

G-Equivariant Multi-Layer Perceptron.

evaluate_output_shape(input_shape: tuple[int, ...]) → tuple[int, ...]

Compute the shape the output tensor which would be generated by this module when a tensor with shape input_shape is provided as input.

Parameters:

input_shape (tuple) – shape of the input tensor

Returns:

shape of the output tensor

export()

Exporting to a torch.nn.Sequential

extra_repr() → str

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x: GeometricTensor) → GeometricTensor

Forward pass of the EMLP.

class symm_learning.models.IMLP(in_type: FieldType, out_dim: int, hidden_units: list[int] = [128, 128, 128], activation: str = 'ReLU', bias: bool = False, hidden_rep: Representation = None)

G-Invariant Multi-Layer Perceptron.

This module is a G-invariant MLP that extracts G-invariant features from the input tensor. The input tensor is first processed by an EMLP module that extracts G-equivariant features. The output of the EMLP module is then processed by an IrrepSubspaceNormPooling module that computes the norm of the features in each G-stable subspace associated to individual irreducible representations. The output of the IrrepSubspaceNormPooling module is a tensor with G-invariant features that can be processed with any NN architecture. Default implementation is to add a single linear layer projecting the invariant features to the desired output dimension.

check_equivariance(atol: float = 1e-06, rtol: float = 0.0001) → list[tuple[any, float]]

Method that automatically tests the equivariance of the current module. The default implementation of this method relies on escnn.nn.GeometricTensor.transform() and uses the the group elements in testing_elements.

This method can be overwritten for custom tests.

Returns:

a list containing containing for each testing element a pair with that element and the corresponding equivariance error

evaluate_output_shape(input_shape: tuple[int, ...]) → tuple[int, ...]

Compute the shape the output tensor which would be generated by this module when a tensor with shape input_shape is provided as input.

Parameters:

input_shape (tuple) – shape of the input tensor

Returns:

shape of the output tensor

export()

Exporting to a torch.nn.Sequential

forward(x: GeometricTensor) → GeometricTensor

Forward pass of the G-invariant MLP.