CNN -> Model

Convolutions Modules

class tardis_em.cnn.model.convolution.GeLU(tanh: float | None = None)

Applies the Gaussian Error Linear Unit (GeLU) activation function.

The GeLU activation function is a smooth approximation to the ReLU activation function. This implementation provides an optional tanh parameter that can be used to scale the input argument to the standard mathematical error function (erf). It is primarily used in neural networks to introduce non-linearity.

forward(x: Tensor) → Tensor

Applies a transformation to the input tensor using the hyperbolic tangent error function, scaling the result accordingly. This method is commonly utilized in machine learning models and functions by smoothing input data values into a specified scaling range.

Parameters:

x (torch.Tensor) – Input tensor on which the transformation is applied.

Returns:

Transformed tensor after applying the scaled error function operation.

Return type:

torch.Tensor

tardis_em.cnn.model.convolution.convolution(in_ch: int, out_ch: int, components: str, kernel: int, padding: int, num_group=None) → list

Builds a neural network block by assembling components specified in the given string. This function enables the construction of customizable CNN layers based on the input parameters and component configurations. The components can consist of various operations such as convolutions (2D or 3D), normalization techniques (GroupNorm, BatchNorm), and activation functions (ReLU, LeakyReLU, GeLU, PReLU). These components are added sequentially to the module list based on their order in the components string.

Parameters:

• in_ch (int) – Number of input channels for the convolutional layers.

• out_ch (int) – Number of output channels for the convolutional layers.

• components (str) – A string that specifies the sequence of operations for the CNN block. Each character or group represents a specific operation such as convolution, normalization, or activation.

• kernel (int or tuple) – Kernel size for the convolution. Can be an integer or a tuple representing the size.

• padding (int or tuple) – Padding size for the convolution. Can be an integer or a tuple representing the size.

• num_group (int, optional) – Number of groups for Group Normalization. Required if GroupNorm is included in the components. Defaults to None.

Returns:

A list of tuples representing the layers in the CNN block. Each tuple contains a layer name and its corresponding module.

Return type:

list

class tardis_em.cnn.model.convolution.SingleConvolution(in_ch: int, out_ch: int, components: str, kernel: int, padding: int, num_group=None, block_type='any')

Represents a sequential layer that contains a single convolution operation.

This class is a part of a neural network building process, specifically for performing either 2D or 3D convolutional operations. It inherits from nn.Sequential to combine multiple convolutional modules into a sequence based on the specified dimensionality of the operation. It dynamically constructs convolutional layers considering the input parameters.

class tardis_em.cnn.model.convolution.DoubleConvolution(in_ch: int, out_ch: int, block_type: str, kernel: int, padding: int, components='cgr', num_group=None)

Implements a double convolutional block feature in a neural network.

This class is a specialized neural network module implementing a double convolutional operation, designed either as an encoder block for reducing spatial resolution while increasing the number of features or as a decoder block for recovering spatial resolution with retained channel properties. It supports two block types - “encoder” and “decoder”, leveraging the flexibility of defining kernel size, padding, and customizable components. The concept of group convolutions is enabled for enhanced modularity via num_group.

class tardis_em.cnn.model.convolution.RecurrentDoubleConvolution(in_ch: int, out_ch: int, block_type: str, kernel: int, padding: int, components='cgr', num_group=None)

Defines the RecurrentDoubleConvolution class, which implements a customized recurrent double convolution block designed for encoder and decoder operations in convolutional neural networks (CNN). The block consists of three consecutive single convolution operations with different configurations. It incorporates residual connections and supports various non-linearity components such as LeakyReLU, ReLU, and PReLU, making it suitable for flexible deep learning architectures.

forward(x: Tensor) → Tensor

Processes input tensor through multiple convolutional layers and applies a non-linearity function. Implements a residual connection between intermediate convolutions and the final output computation.

Parameters:

x (torch.Tensor) – Input tensor to process through the defined convolutional layers.

Returns:

Transformed tensor after applying convolutional layers, residual connection, and non-linearity.

Return type:

torch.Tensor

Generalize Encoder Module

class tardis_em.cnn.model.encoder_blocks.EncoderBlock(in_ch: int, out_ch: int, conv_module, conv_kernel=3, max_pool=True, dropout: float | None = None, pool_kernel=2, padding=1, components='3gcr', num_group=8, attn_features=False)

Represents an encoder block with convolutional modules, optional max-pooling, dropout, and attention features for deep learning architectures.

This class constructs a configurable encoder block. The encoder block supports max-pooling layers, dropout layers, attention features, and employs convolutional modules with specific kernel sizes and padding. The components used in the block can be customized through component identifiers.

dropout

Optionally, add maxpool

forward(x: Tensor) → Tensor

Processes input tensor through layers including convolutional, attention mechanisms, optional pooling, and dropout, returning the transformed tensor. This method serves as the key forward computation logic for the module.

Parameters:

x (torch.Tensor) – Input tensor to be processed.

Returns:

Processed output tensor.

Return type:

torch.Tensor

tardis_em.cnn.model.encoder_blocks.build_encoder(in_ch: int, conv_layers: int, conv_layer_scaler: int, conv_kernel: int, padding: int, num_group: int, components: str, pool_kernel: int, conv_module, dropout: float | None = None, attn_features=False) → ModuleList

Constructs and returns a sequence of encoder blocks as a module list. Each encoder block is defined based on the provided parameters and forms a hierarchical structure of feature extraction layers. The first encoder block does not use max pooling, while subsequent layers include max pooling operations as specified. The method also accommodates advanced features such as attention mechanisms and grouped convolutions.

Parameters:

• in_ch – Number of input image channels for the first encoder block. Each subsequent block adapts based on feature scaling.

• conv_layers – Number of convolutional layers to create in the encoder architecture.

• conv_layer_scaler – Multiplier to compute the number of feature maps at each level of the encoder.

• conv_kernel – Kernel size for convolution operations in the encoder blocks. Can be an integer or tuple.

• padding – Padding value for the convolutional layers. Accepts integer or tuple for specific layer configurations.

• num_group – Number of groups for grouped convolution operations.

• components – String identifier or configuration for additional encoder components such as normalization layers or activation functions.

• pool_kernel – Kernel size used for max pooling operations across encoder blocks except the first block. Accepts integer or tuple.

• conv_module – A callable or module defining the specific implementation for convolutional operations.

• dropout – Optional dropout probability to apply regularization in each encoder block. Defaults to None, disabling dropout.

• attn_features – Boolean flag to enable or disable attention mechanisms in the encoder blocks. Defaults to False.

Returns:

Returns a nn.ModuleList containing the constructed encoder blocks.

Generalize Decoder Module

class tardis_em.cnn.model.decoder_blocks.DecoderBlockCNN(in_ch: int, out_ch: int, conv_kernel: int, padding: int, size: int, dropout: float | None = None, components='3gcr', num_group=8)

DecoderBlockCNN implements a CNN-based decoder block architecture suitable for use in U-Net-like neural networks. It enables upscaling and optional dropout mechanism for regularization.

This class builds a modular decoder block consisting of upscaling (using either bilinear or trilinear interpolation), a double convolution module configured for decoding tasks, and optional dropout. The class also permits initialization of the weights for its layers.

dropout

Build decoders

deconv_module

Optional Dropout

forward(encoder_features: Tensor, x: Tensor) → Tensor

Combines encoder features and upscale features, applies a deconvolution module, and optionally applies a dropout layer to the output during the forward pass of the network.

Parameters:

• encoder_features – Input tensor representing features extracted from the encoder block.

• x – Input tensor that has been upsampled by the current block of the decoder.

Returns:

A tensor obtained by combining the encoder features and the upsampled tensor, after applying deconvolution and optional dropout.

class tardis_em.cnn.model.decoder_blocks.DecoderBlockRCNN(in_ch: int, out_ch: int, conv_kernel: int, padding: int, size: int, dropout: float | None = None, components='3gcr', num_group=8)

Implements a Decoder Block for a Recurrent Convolutional Neural Network (RCNN).

This class defines a decoder block structure for RCNN models. The decoder block performs upscaling, residual connections, convolutional operations, and optional dropout regularization. It is designed for flexible integration into various neural network architectures, particularly those requiring upscaling and deep feature convolutions.

dropout

Build decoders

deconv_res_module

Optional Dropout

forward(encoder_features: Tensor, x: Tensor) → Tensor

The forward method implements the forward pass in the neural network layer. It modifies input using an upscale operation, applies a deconvolution module, and adds the result with encoder features. It further processes the data with a deconvolution residual module. If a dropout is specified, it applies the dropout layer to the processed output, and finally returns the result.

Parameters:

• encoder_features – Features derived from the encoder network to be combined with the processed tensor.

• x – Input tensor passed through this forward method for processing.

Returns:

The processed tensor after upscaling, combination with encoder features, residual operations, and optional dropout application.

class tardis_em.cnn.model.decoder_blocks.DecoderBlockUnet3Plus(in_ch: int, out_ch: int, conv_kernel: int, padding: int, size: int, components: str, num_group: int, num_layer: int, encoder_feature_ch: list, attn_features=False, dropout: float | None = None)

DecoderBlockUnet3Plus is a class for decoding layers in a UNet3+ architecture. It is designed for multi-scale feature fusion with optional attention mechanism, skip connections, and dropout capabilities. This block helps upscale encoded spatial representations and recombine them with corresponding encoder features.

dropout

Main Block Up-Convolution

deconv

Skip-Connection Encoders

forward(x: Tensor, encoder_features: list) → Tensor

Processes the input tensor through the forward pass of a neural network block.

This function handles different operations, including upscaling, deconvolution, skip-connections with encoder features, optional attention mechanisms, and dropout application. It integrates different components such as upscaling the input tensor, merging encoder features, applying optional attention layers, and optionally applying dropout operation at the end.

Parameters:

• x (torch.Tensor) – Input tensor to be processed through the block.

• encoder_features (list[torch.Tensor]) – List of tensors representing the outputs from encoder layers.

Returns:

Processed tensor obtained after applying all operations in the block.

Return type:

torch.Tensor

tardis_em.cnn.model.decoder_blocks.build_decoder(conv_layers: int, conv_layer_scaler: int, components: str, num_group: int, conv_kernel: int, padding: int, sizes: list, dropout: float | None = None, deconv_module='CNN', attn_features=False)

Builds a decoder module consisting of multiple decoder blocks. The choice of decoder block type depends on the deconv_module parameter, which supports CNN, RCNN, and Unet3Plus decoder types. The decoder is constructed based on the feature map channels derived from the number of convolution layers and a scaling factor. Each decoder block is appended to a list, which is eventually returned as a ModuleList.

Parameters:

• conv_layers – Number of convolution layers to base the feature computation.

• conv_layer_scaler – Scaling factor to adjust the number of features for each layer.

• components – Type or characteristic of components used in the decoder blocks.

• num_group – Number of groups for grouped convolutions, if applicable.

• conv_kernel – Size of the convolution kernel used within the decoder.

• padding – Padding to be applied in the convolution layers.

• sizes – List of spatial sizes for each decoder block output.

• dropout – Fraction for dropout regularization to prevent overfitting (optional).

• deconv_module – Type of the decoder module, options are “CNN”, “RCNN”, or “unet3plus”.

• attn_features – Flag to include attention features for Unet3Plus decoder (specifically when deconv_module=”unet3plus”).

Returns:

A torch.nn.ModuleList containing the decoder blocks with the configured parameters.

Weight Initialization

tardis_em.cnn.model.init_weights.weights_init_kaiming(m)

Initializes the weights of layers in a neural network module using the Kaiming initialization technique. This function checks the type of layer contained in the past module and applies different initialization strategies depending on the layer type. Specifically, it adjusts the weights and biases for convolutional layers, batch normalization layers, and group normalization layers. This function is commonly used for improving the convergence of deep learning models with ReLU activation functions.

Parameters:

m – The module whose weights need to be initialized. It is expected to be an instance of a layer or a module class such as Conv2d, Conv3d, BatchNorm, or GroupNorm.

Returns:

None

tardis_em.cnn.model.init_weights.init_weights(net, init_type='kaiming')

Initializes the weights of a neural network based on the specified initialization type. The function applies the kaiming initialization if init_type is set to “kaiming”. If the provided init_type is not implemented, an error is raised.

Parameters:

• net – The neural network whose weights need to be initialized.

• init_type – A string specifying the type of weight initialization to be used. Defaults to “kaiming”. Must be one of the implemented initialization methods.

Returns:

None