Skip to content

SCNN

This module defines a Steerable Convolutional Neural Network (SteerableCNN) using the escnn library for equivariant neural networks. The primary class, SteerableCNN, allows for the creation of a convolutional neural network that handles various symmetries, making it useful for tasks requiring rotational invariance.

The module imports necessary libraries, including torch and escnn, and defines utility functions such as create_block, compute_new_features, create_pool, and create_final_block. These functions help construct convolutional blocks, calculate feature map sizes, create anti-aliased pooling layers, and build fully connected layers.

The SteerableCNN class initializes with a configuration dictionary and a symmetry parameter, setting up parameters like kernel sizes, paddings, strides, and feature fields. It determines the symmetry group and initializes the input type for the network.

The network is built by iterating through specified kernel sizes, creating convolutional blocks and pooling layers, adding a group pooling layer for invariance, and creating final fully connected layers using create_final_block.

The forward method processes the input tensor through the network, applying each equivariant block, performing group pooling, and classifying the output using the fully connected layers.

Overall, this module provides a flexible framework for building steerable convolutional neural networks with configurable symmetries and architectures.

create_final_block(config, n_initial, nscreens)

Creates a fully connected neural network block based on a predefined configuration.

This function dynamically creates a sequence of PyTorch layers for a fully connected neural network. The configuration for the layers is read from a global config dictionary which should contain a 'final_block' key with a list of layer configurations. Each layer configuration is a dictionary that must include a 'type' key with the name of the layer class (e.g., 'Linear', 'Dropout', etc.) and can include additional keys for the layer parameters.

The first 'Linear' layer in the configuration has its number of input features set to n_in, and any 'Linear' layer with 'out_features' set to 'nscreens' has its number of output features set to nscreens.

Args:

Parameters:

  • config (dict) –

    The global configuration dictionary containing the layer configurations.

  • n_initial (int) –

    The number of input features for the first 'Linear' layer.

  • nscreens (int) –

    The number of output features for any 'Linear' layer with 'out_features' set to 'nscreens'.

Returns:

  • torch.nn.Sequential: A sequential container of the configured PyTorch layers. –
Source code in src/speckcn2/scnn.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def create_final_block(config: dict, n_initial: int,
                       nscreens: int) -> nn.Sequential:
    """Creates a fully connected neural network block based on a predefined
    configuration.

    This function dynamically creates a sequence of PyTorch layers for a fully connected
    neural network. The configuration for the layers is read from a global `config` dictionary
    which should contain a 'final_block' key with a list of layer configurations. Each layer
    configuration is a dictionary that must include a 'type' key with the name of the layer
    class (e.g., 'Linear', 'Dropout', etc.) and can include additional keys for the layer
    parameters.

    The first 'Linear' layer in the configuration has its number of input features set to `n_in`,
    and any 'Linear' layer with 'out_features' set to 'nscreens' has its number of output features
    set to `nscreens`.

    Args:
    Parameters
    ----------
    config : dict
        The global configuration dictionary containing the layer configurations.
    n_initial : int
        The number of input features for the first 'Linear' layer.
    nscreens : int
        The number of output features for any 'Linear' layer with 'out_features' set to 'nscreens'.

    Returns
    ----------
    torch.nn.Sequential: A sequential container of the configured PyTorch layers.
    """
    layers = []

    for layer_config in config['final_block']:
        layer_type = layer_config.pop('type')

        if layer_type == 'Linear':
            # The first linear layer has a constrained number of input features
            if n_initial != -1:
                layer_config['in_features'] = n_initial
                n_initial = -1
            if layer_config['out_features'] == 'nscreens':
                layer_config['out_features'] = nscreens
            # Pass the number of features as args
            n_in = layer_config.pop('in_features')
            n_out = layer_config.pop('out_features')

            layers.append(torch.nn.Linear(n_in, n_out, **layer_config))
        else:
            layer_class = getattr(torch.nn, layer_type)
            layers.append(layer_class(**layer_config))

    return torch.nn.Sequential(*layers)