models
A collection of equivariant neural network architectures.
- class symm_learning.models.EMLP(in_type: FieldType, out_type: FieldType, hidden_units: list[int], activation: str | list[str] = 'ReLU', batch_norm: bool = False, pointwise_activation: bool = True, bias: bool = True, hidden_rep: Representation | None = None)[source]
G-Equivariant Multi-Layer Perceptron.
- evaluate_output_shape(input_shape: tuple[int, ...]) tuple[int, ...][source]
Compute the shape the output tensor which would be generated by this module when a tensor with shape
input_shapeis provided as input.- Parameters:
input_shape (tuple) – shape of the input tensor
- Returns:
shape of the output tensor
- 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)[source]
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]][source]
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 intesting_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, ...][source]
Compute the shape the output tensor which would be generated by this module when a tensor with shape
input_shapeis provided as input.- Parameters:
input_shape (tuple) – shape of the input tensor
- Returns:
shape of the output tensor
- class symm_learning.models.MLP(in_dim: int, out_dim: int, hidden_units: list[int], activation: Module | list[Module] = ReLU(), batch_norm: bool = False, bias: bool = True)[source]
Standard baseline MLP. Symmetry-related parameters are ignored.
- class symm_learning.models.TimeCNNEncoder(in_dim: int, out_dim: int, hidden_channels: list[int], horizon: int, activation: Module | list[Module] = ReLU(), batch_norm: bool = False, bias: bool = True, mlp_hidden: list[int] = [128], downsample: str = 'stride', append_last_frame: bool = False)[source]
1D CNN encoder for inputs of shape (N, in_dim, H).
The model applies a stack of 1D conv blocks over time. Each block halves the time horizon using either stride-2 convolution or max-pooling. The final feature map of shape (N, C_L, H_out) is flattened and passed to a 1-hidden-layer MLP head.
Convolutional features are time-translation equivariant; flattening preserves this up to a fixed permutation under shifts. The optional concatenation of the last input frame and the final MLP can break output equivariance.
- Parameters:
in_dim – Input channel dimension.
out_dim – Output feature dimension.
hidden_channels – Channels per conv block; depth equals len(hidden_channels).
horizon – Input sequence length H.
activation – Activation module or list (one per block). If a single module is given, it is replicated for all blocks.
batch_norm – If True, add BatchNorm1d after each convolution.
bias – Use bias in conv/linear layers.
mlp_hidden – Hidden units of the final MLP head (list for deeper heads).
downsample – “stride” (default) or “pooling”; each block halves H.
append_last_frame – If True, concatenate x[:, :, -1] (N, in_dim) to flattened conv features before the MLP.
- Returns:
Tensor of shape (N, out_dim).
- class symm_learning.models.eTimeCNNEncoder(in_type: FieldType, out_type: FieldType, hidden_channels: list[int], time_horizon: int, activation: Module | list[Module] = ReLU(), batch_norm: bool = False, bias: bool = True, mlp_hidden: list[int] = [128], downsample: str = 'stride', append_last_frame: bool = False)[source]
Equivariant 1D CNN encoder using eConv1D and eBatchNorm1d.
Processes inputs of shape (N, in_dim, H) through L equivariant conv blocks (stride-2 only). The flattened feature map is fed to an equivariant (EMLP) or invariant (IMLP) head depending on the requested out_type. Pooling is not supported in the equivariant blocks.
- Parameters:
in_type – Input FieldType defining channel representation structure.
out_type – Desired output field type. If it contains only trivial irreps, an invariant head (IMLP) is used; otherwise an equivariant head (EMLP) is used.
hidden_channels – Channels per conv block; depth equals len(hidden_channels).
horizon – Input sequence length H.
activation – Activation module or list (one per block). Only pointwise activations are supported equivariantly in the blocks. ReLU/ELU/LeakyReLU/Mish are mapped; others fall back to PointwiseNonLinearity with a matching name.
batch_norm – If True, add eBatchNorm1d after each convolution.
bias – Use bias in conv/linear layers.
mlp_hidden – Hidden units of the head MLP as a list (single hidden layer by default).
downsample – Must be “stride”. If “pooling” is requested, NotImplementedError is raised.
append_last_frame – If True, concatenate the last input frame (N, in_type.size) to the flattened conv features before the head. This can break output equivariance.
- Returns:
GeometricTensor with field type out_type.