equiadapt.images.canonicalization_networks package

Submodules

equiadapt.images.canonicalization_networks.custom_equivariant_networks module

class equiadapt.images.canonicalization_networks.custom_equivariant_networks.CustomEquivariantNetwork(in_shape: Tuple[int, int, int, int], out_channels: int, kernel_size: int, group_type: str = 'rotation', num_rotations: int = 4, num_layers: int = 1, device: str = 'cpu')

Bases: Module

This class represents a custom equivariant network.

The network is equivariant to a specified group, which can be either the rotation group or the roto-reflection group. The network consists of a sequence of equivariant convolutional layers, each followed by a ReLU activation function.

__init__()

Initializes the CustomEquivariantNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, group_size).

Return type:

torch.Tensor

equiadapt.images.canonicalization_networks.custom_group_equivariant_layers module

class equiadapt.images.canonicalization_networks.custom_group_equivariant_layers.RotationEquivariantConv(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a rotation equivariant convolutional layer.

The layer is equivariant to a group of rotations. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotationEquivariantConv instance.

get_rotated_permuted_weights()

Returns the weights of the layer after rotation and permutation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, num_rotations, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_rotations, height, width).

Return type:

torch.Tensor

get_rotated_permuted_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation and permutation.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation and permutation.

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.custom_group_equivariant_layers.RotationEquivariantConvLift(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a rotation equivariant convolutional layer with lifting.

The layer is equivariant to a group of rotations. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotationEquivariantConvLift instance.

get_rotated_weights()

Returns the weights of the layer after rotation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_rotations, height, width).

Return type:

torch.Tensor

get_rotated_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation.

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.custom_group_equivariant_layers.RotoReflectionEquivariantConv(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a roto-reflection equivariant convolutional layer.

The layer is equivariant to a group of rotations and reflections. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotoReflectionEquivariantConv instance.

get_rotoreflected_permuted_weights()

Returns the weights of the layer after rotation, reflection, and permutation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, num_group_elements, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_group_elements, height, width).

Return type:

torch.Tensor

get_rotoreflected_permuted_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation, reflection, and permutation.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation, reflection, and permutation.

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.custom_group_equivariant_layers.RotoReflectionEquivariantConvLift(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a roto-reflection equivariant convolutional layer with lifting.

The layer is equivariant to a group of rotations and reflections. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotoReflectionEquivariantConvLift instance.

get_rotoreflected_weights()

Returns the weights of the layer after rotation, reflection, and permutation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_group_elements, height, width).

Return type:

torch.Tensor

get_rotoreflected_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation and reflection.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation, reflection, and permutation.

Return type:

torch.Tensor

equiadapt.images.canonicalization_networks.custom_nonequivariant_networks module

class equiadapt.images.canonicalization_networks.custom_nonequivariant_networks.ConvNetwork(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a convolutional neural network.

The network consists of a sequence of convolutional layers, each followed by batch normalization and a GELU activation function. The number of output channels of the convolutional layers increases after every third layer. The network ends with a fully connected layer.

__init__()

Initializes the ConvNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, out_vector_size).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.custom_nonequivariant_networks.ResNet18Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a neural network based on the ResNet-18 architecture.

The network uses a pre-trained ResNet-18 model without its weights. The final fully connected layer of the ResNet-18 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.custom_nonequivariant_networks.WideResNet101Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a neural network based on the WideResNetNetwork architecture.

The network uses a pre-trained WideResNet model. The final fully connected layer of the WideResNet101 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.custom_nonequivariant_networks.WideResNet50Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a neural network based on the WideResNetNetwork architecture.

The network uses a pre-trained WideResNet model. The final fully connected layer of the WideResNet50 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

equiadapt.images.canonicalization_networks.escnn_networks module

class equiadapt.images.canonicalization_networks.escnn_networks.ESCNNEquivariantNetwork(in_shape: tuple, out_channels: int, kernel_size: int, group_type: str = 'rotation', num_rotations: int = 4, num_layers: int = 1)

Bases: Module

This class represents an Equivariant Convolutional Neural Network (Equivariant CNN).

The network is equivariant to a group of transformations, which can be either rotations or roto-reflections. The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization, a ReLU activation function, and dropout. The number of output channels of the convolutional layers is the same for all layers.

__init__()

Initializes the ESCNNEquivariantNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, num_group_elements).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.escnn_networks.ESCNNSteerableNetwork(in_shape: tuple, out_channels: int, kernel_size: int = 9, group_type: str = 'rotation', num_layers: int = 1)

Bases: Module

This class represents a Steerable Equivariant Convolutional Neural Network (Equivariant CNN).

The network is equivariant under all planar rotations. The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a FourierELU activation function. The number of output channels of the convolutional layers is the same for all layers.

__init__()

Initializes the ESCNNSteerableNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 2, 2).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.escnn_networks.ESCNNWRNEquivariantNetwork(in_shape: tuple, out_channels: int = 64, kernel_size: int = 9, group_type: str = 'rotation', num_layers: int = 12, num_rotations: int = 4)

Bases: Module

This class represents a Wide Residual Network (WRN) that is equivariant under rotations or roto-reflections.

The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()

Initializes the ESCNNWRNEquivariantNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, group_size).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.escnn_networks.ESCNNWideBasic(in_type: FieldType, middle_type: FieldType, out_type: FieldType, kernel_size: int = 3)

Bases: EquivariantModule

This class represents a wide basic layer for an Equivariant Convolutional Neural Network (Equivariant CNN).

The layer consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()

Initializes the ESCNNWideBasic instance.

forward()

Performs a forward pass through the layer.

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

forward(x: GeometricTensor) → GeometricTensor

Performs a forward pass through the layer.

Parameters:

x (e2cnn.nn.GeometricTensor) – The input data.

Returns:

The output of the layer. The input is added to the output (residual connection).

Return type:

e2cnn.nn.GeometricTensor

class equiadapt.images.canonicalization_networks.escnn_networks.ESCNNWideBottleneck(in_type: FieldType, middle_type: FieldType, out_type: FieldType, kernel_size: int = 3)

Bases: EquivariantModule

This class represents a wide bottleneck layer for an Equivariant Convolutional Neural Network (Equivariant CNN).

The layer consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()

Initializes the ESCNNWideBottleneck instance.

forward()

Performs a forward pass through the layer.

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

forward(x: GeometricTensor) → GeometricTensor

Performs a forward pass through the layer.

Parameters:

x (e2cnn.nn.GeometricTensor) – The input data.

Returns:

The output of the layer. The input is added to the output (residual connection).

Return type:

e2cnn.nn.GeometricTensor

Module contents

class equiadapt.images.canonicalization_networks.ConvNetwork(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a convolutional neural network.

The network consists of a sequence of convolutional layers, each followed by batch normalization and a GELU activation function. The number of output channels of the convolutional layers increases after every third layer. The network ends with a fully connected layer.

__init__()

Initializes the ConvNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, out_vector_size).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.CustomEquivariantNetwork(in_shape: Tuple[int, int, int, int], out_channels: int, kernel_size: int, group_type: str = 'rotation', num_rotations: int = 4, num_layers: int = 1, device: str = 'cpu')

Bases: Module

This class represents a custom equivariant network.

The network is equivariant to a specified group, which can be either the rotation group or the roto-reflection group. The network consists of a sequence of equivariant convolutional layers, each followed by a ReLU activation function.

__init__()

Initializes the CustomEquivariantNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, group_size).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.ESCNNEquivariantNetwork(in_shape: tuple, out_channels: int, kernel_size: int, group_type: str = 'rotation', num_rotations: int = 4, num_layers: int = 1)

Bases: Module

This class represents an Equivariant Convolutional Neural Network (Equivariant CNN).

The network is equivariant to a group of transformations, which can be either rotations or roto-reflections. The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization, a ReLU activation function, and dropout. The number of output channels of the convolutional layers is the same for all layers.

__init__()

Initializes the ESCNNEquivariantNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, num_group_elements).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.ESCNNSteerableNetwork(in_shape: tuple, out_channels: int, kernel_size: int = 9, group_type: str = 'rotation', num_layers: int = 1)

Bases: Module

This class represents a Steerable Equivariant Convolutional Neural Network (Equivariant CNN).

The network is equivariant under all planar rotations. The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a FourierELU activation function. The number of output channels of the convolutional layers is the same for all layers.

__init__()

Initializes the ESCNNSteerableNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 2, 2).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.ESCNNWRNEquivariantNetwork(in_shape: tuple, out_channels: int = 64, kernel_size: int = 9, group_type: str = 'rotation', num_layers: int = 12, num_rotations: int = 4)

Bases: Module

This class represents a Wide Residual Network (WRN) that is equivariant under rotations or roto-reflections.

The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()

Initializes the ESCNNWRNEquivariantNetwork instance.

forward()

Performs a forward pass through the network.

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, group_size).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.ESCNNWideBasic(in_type: FieldType, middle_type: FieldType, out_type: FieldType, kernel_size: int = 3)

Bases: EquivariantModule

This class represents a wide basic layer for an Equivariant Convolutional Neural Network (Equivariant CNN).

The layer consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()

Initializes the ESCNNWideBasic instance.

forward()

Performs a forward pass through the layer.

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

forward(x: GeometricTensor) → GeometricTensor

Performs a forward pass through the layer.

Parameters:

x (e2cnn.nn.GeometricTensor) – The input data.

Returns:

The output of the layer. The input is added to the output (residual connection).

Return type:

e2cnn.nn.GeometricTensor

class equiadapt.images.canonicalization_networks.ESCNNWideBottleneck(in_type: FieldType, middle_type: FieldType, out_type: FieldType, kernel_size: int = 3)

Bases: EquivariantModule

This class represents a wide bottleneck layer for an Equivariant Convolutional Neural Network (Equivariant CNN).

The layer consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()

Initializes the ESCNNWideBottleneck instance.

forward()

Performs a forward pass through the layer.

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

forward(x: GeometricTensor) → GeometricTensor

Performs a forward pass through the layer.

Parameters:

x (e2cnn.nn.GeometricTensor) – The input data.

Returns:

The output of the layer. The input is added to the output (residual connection).

Return type:

e2cnn.nn.GeometricTensor

class equiadapt.images.canonicalization_networks.ResNet18Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a neural network based on the ResNet-18 architecture.

The network uses a pre-trained ResNet-18 model without its weights. The final fully connected layer of the ResNet-18 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.RotationEquivariantConv(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a rotation equivariant convolutional layer.

The layer is equivariant to a group of rotations. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotationEquivariantConv instance.

get_rotated_permuted_weights()

Returns the weights of the layer after rotation and permutation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, num_rotations, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_rotations, height, width).

Return type:

torch.Tensor

get_rotated_permuted_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation and permutation.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation and permutation.

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.RotationEquivariantConvLift(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a rotation equivariant convolutional layer with lifting.

The layer is equivariant to a group of rotations. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotationEquivariantConvLift instance.

get_rotated_weights()

Returns the weights of the layer after rotation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_rotations, height, width).

Return type:

torch.Tensor

get_rotated_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation.

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.RotoReflectionEquivariantConv(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a roto-reflection equivariant convolutional layer.

The layer is equivariant to a group of rotations and reflections. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotoReflectionEquivariantConv instance.

get_rotoreflected_permuted_weights()

Returns the weights of the layer after rotation, reflection, and permutation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, num_group_elements, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_group_elements, height, width).

Return type:

torch.Tensor

get_rotoreflected_permuted_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation, reflection, and permutation.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation, reflection, and permutation.

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.RotoReflectionEquivariantConvLift(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')

Bases: Module

This class represents a roto-reflection equivariant convolutional layer with lifting.

The layer is equivariant to a group of rotations and reflections. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()

Initializes the RotoReflectionEquivariantConvLift instance.

get_rotoreflected_weights()

Returns the weights of the layer after rotation, reflection, and permutation.

forward()

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_group_elements, height, width).

Return type:

torch.Tensor

get_rotoreflected_weights(weights: Tensor, num_rotations: int = 4) → Tensor

Returns the weights of the layer after rotation and reflection.

Parameters:

• weights (torch.Tensor) – The weights of the layer.

• num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation, reflection, and permutation.

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.WideResNet101Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a neural network based on the WideResNetNetwork architecture.

The network uses a pre-trained WideResNet model. The final fully connected layer of the WideResNet101 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

class equiadapt.images.canonicalization_networks.WideResNet50Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)

Bases: Module

This class represents a neural network based on the WideResNetNetwork architecture.

The network uses a pre-trained WideResNet model. The final fully connected layer of the WideResNet50 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor