CNN -> Model

Convolutions Modules

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

CUSTOM GAUSSIAN ERROR LINEAR UNITS ACTIVATION FUNCTION

Parameters:

tanh (float) – hyperbolic tangent value for GeLU

forward(x: Tensor) → Tensor

Forward GeLu transformation.

Parameters:

x (torch.Tensor) – Image tensor to transform.

Returns:

Transformed image tensor.

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

Customizable convolution block builder.

Build an convolution block with a specified components and order: dimension (2 or 3), conv (c), ReLu (r), LeakyReLu (l), GeLu (e), PReLu (p), GroupNorm (g), BatchNorm (b).

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • components (str) – Components that are used for conv. block.

  • kernel (int, tuple) – Kernel size for the convolution.

  • padding (int, tuple) – Padding size for the convolution.

  • num_group (int) – Num. of groups for the nn.GroupNorm. None -> if nn.GroupNorm is not used.

Returns:

Ordered nn.Module list.

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')

STANDARD SINGLE 3D CONVOLUTION BLOCK

Output single convolution composed of conv, normalization and relu in order defined by components variable.

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • block_type (str) – Define encode or decoder path e.g. - ‘encoder’: Encoder convolution path - ‘decoder’: Decoder convolution path

  • components (str) – Components that are used for conv. block.

  • kernel (int, tuple) – Kernel size for the convolution.

  • padding (int, tuple) – Padding size for the convolution.

  • num_group (int) – Num. of groups for the nn.GroupNorm. None -> if nn.GroupNorm is not used.

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)

DOUBLE CONVOLUTION BLOCK

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • block_type (str) – Define encode or decoder path e.g. - ‘encoder’: Encoder convolution path - ‘decoder’: Decoder convolution path

  • components (str) – Components that are used for conv. block.

  • kernel (int, tuple) – Kernel size for the convolution.

  • padding (int, tuple) – Padding size for the convolution.

  • num_group (int) – Num. of groups for the nn.GroupNorm. None -> if nn.GroupNorm is not used.

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)

RECURRENT DOUBLE CONVOLUTION BLOCK

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • block_type (str) – Define encode or decoder path e.g. - ‘encoder’: Encoder convolution path - ‘decoder’: Decoder convolution path

  • components (str) – Components that are used for conv. block.

  • kernel (int, tuple) – Kernel size for the convolution.

  • padding (int, tuple) – Padding size for the convolution.

  • num_group (int) – Num. of groups for the nn.GroupNorm. None -> if nn.GroupNorm is not used.

forward(x: Tensor) → Tensor

Forward function for customized convolution :param x: Image patch.

Returns:

Up or down convoluted image patch.

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)

ENCODER BUILDER

Single encoder module composed of nn.MaxPool and convolution module.

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • conv_module (conv_module) – Single, Double or RCNN convolution block.

  • conv_kernel (int) – Convolution kernel size.

  • max_pool (int) – If True nn.MaxPool is applied.

  • pool_kernel (int) – Kernel size for max pooling.

  • dropout (float, optional) – Optionals, dropout rate.

  • padding (int) – Padding size for the convolution.

  • components (str) – Components that are used for conv. block.

  • num_group (int) – Num. of groups for the nn.GroupNorm. None -> if nn.GroupNorm is not used.

dropout

Optionally, add maxpool

forward(x: Tensor) → Tensor

Forward CNN encoder block.

Parameters:

x (torch.Tensor) – Image torch before convolution.

Returns:

Image after convolution.

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

Encoder wrapper for entire CNN model.

Create encoder block from feature map and convolution modules. Number of encoder layers is indicated by number of features.

Parameters:

  • in_ch (int) – Number of input channels.

  • conv_layers (int) – Number of convolution layers.

  • conv_layer_scaler (int) – Number of channel by which each CNN block is scaled up.

  • conv_module (conv_module) – Single, Double or RCNN convolution block.

  • conv_kernel (int) – Convolution kernel size.

  • pool_kernel (int) – Kernel size for max pooling.

  • dropout (float, optional) – Optionals, dropout rate.

  • padding (int) – Padding size for the convolution.

  • components (str) – Components that are used for conv. block.

  • num_group (int) – Num. of groups for the nn.GroupNorm. None -> if nn.GroupNorm is not used.

  • attn_features (bool)

Returns:

Encoder block.

Return type:

nn.ModuleList

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)

CNN DECODER BUILDER FOR UNET

Create a decoder block consisting of indicated number of deconvolutions followed by up-sampling and concatenation.

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • size (int) – Size for the resampling.

  • dropout (float, optional) – Optional, dropout rate.

  • components (str) – String list of components from which convolution block is composed

  • num_group (int) – No. of groups for nn.GroupNorm()

dropout

Build decoders

deconv_module

Optional Dropout

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

Forward CNN decoder block for Unet

Parameters:

  • encoder_features (torch.Tensor) – Residual connection.

  • x (torch.Tensor) – Image tensor before convolution.

Returns:

Image tensor after convolution.

Return type:

torch.Tensor

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)

RCNN DECODER BUILDER

Create a decoder block consisting of indicated number of deconvolutions followed by upscale and connection with torch.cat().

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • size (int) – Size for the resampling.

  • dropout (float, optional) – Optional, dropout rate.

  • components (str) – String list of components from which convolution block is composed

  • num_group (int) – No. of groups for nn.GroupNorm()

dropout

Build decoders

deconv_res_module

Optional Dropout

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

Forward RCNN decoder block

Parameters:

  • encoder_features (torch.Tensor) – Residual connection.

  • x (torch.Tensor) – Image tensor before convolution.

Returns:

Image tensor after convolution.

Return type:

torch.Tensor

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)

CNN DECODER BUILDER FOR UNET3PLUS

Create a decoder block consisting of the indicated number of deconvolutions followed by upscale and connection with torch.cat().

Parameters:

  • in_ch (int) – Number of input channels.

  • out_ch (int) – Number of output channels.

  • size (int) – Size for the resampling.

  • dropout (float, optional) – Optional, dropout rate.

  • components (str) – String list of components from which convolution block is composed.

  • num_group (int) – Number of groups inc CNN block.

  • num_layer (int) – Number of CNN layers.

  • encoder_feature_ch (list) – List of encode outputs.

  • num_group – No. of groups for nn.GroupNorm()

dropout

Main Block Up-Convolution

deconv

Skip-Connection Encoders

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

Forward CNN decoder block for Unet3Plus

Parameters:

  • x (torch.Tensor) – Image tensor before convolution.

  • encoder_features (list) – Residual connection from the encoder.

Returns:

Image tensor after convolution.

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)

Decoder wrapper for the entire CNN model.

Create a decoder block from a number of convolution and convolution modules. The Decoder block is followed by upscale(interpolation) and joining with torch.cat()

Parameters:

  • conv_layers (int) – Number of deconvolution layers.

  • conv_layer_scaler (int) – Number of channels by which each CNN block is scaled up.

  • components (str) – Components that are used for deconvolution block.

  • conv_kernel (int) – Convolution kernel size.

  • padding (int) – Padding size for the convolution.

  • sizes (list) – List of tensor sizes for upscale.

  • num_group (int) – Num. of groups for the nn.GroupNorm. None -> if nn.GroupNorm is not used.

  • dropout (float, Optional) – Dropout value.

  • deconv_module – Module of the deconvolution for the decoder.

  • attn_features

Returns:

List of decoder blocks.

Return type:

nn.ModuleList

Weight Initialization

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

Kaiming weight and bias initialization.

Parameters:

m – CNN block.

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

Wrapper for network module initialization.

Parameters:

  • net – Network to initialized.

  • init_type (str) – type of initialization.

Raises:

NotImplementedError – _description_