connectomics.model

Architectures

class connectomics.model.arch.DeepLabV3(name, backbone_type, out_channel=1, aux_out=False, **kwargs)

Implements DeepLabV3 model from “Rethinking Atrous Convolution for Semantic Image Segmentation”. This implementation only supports 2D inputs. Pretrained ResNet weights on the ImgeaNet dataset is loaded by default.

Parameters

• backbone (nn.Module) – the network used to compute the features for the model. The backbone should return an OrderedDict[Tensor], with the key being “out” for the last feature map used, and “aux” if an auxiliary classifier is used.

• classifier (nn.Module) – module that takes the “out” element returned from the backbone and returns a dense prediction.

• aux_classifier (nn.Module, optional) – auxiliary classifier used during training

• name (str) –

• backbone_type (str) –

• out_channel (int) –

• aux_out (bool) –

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.Discriminator3D(in_channel=1, filters=[64, 64, 128, 128, 256], pad_mode='replicate', act_mode='leaky_relu', norm_mode='in', dilation=1, is_isotropic=False, isotropy=[False, False, False, True, True], stride_list=[2, 2, 2, 2, 1])

3D PatchGAN discriminator

Parameters

• in_channel (int) – number of input channels. Default: 1

• filters (List[int]) – number of filters at each U-Net stage. Default: [32, 64, 96, 96, 96]

• pad_mode (str) – one of 'zeros', 'reflect', 'replicate' or 'circular'. Default: ``’replicate’`

• act_mode (str) – one of 'relu', 'leaky_relu', 'elu', 'gelu', 'swish', 'efficient_swish' or 'none'. Default: 'elu'

• norm_mode (str) – one of 'bn', 'sync_bn' 'in' or 'gn'. Default: 'in'

• dilation (int) – dilation rate of the conv kernels. Default: 1

• is_isotropic (bool) – whether the whole model is isotropic. Default: False

• isotropy (List[bool]) – specify each discriminator layer is isotropic or anisotropic. All elements will be True if is_isotropic is True. Default: [False, False, False, True, True]

• stride_list (List[int]) – list of strides for each conv layer. Default: [2, 2, 2, 2, 1]

Return type

None

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.FPN3D(backbone_type='resnet', block_type='residual', feature_keys=['feat1', 'feat2', 'feat3', 'feat4', 'feat5'], in_channel=1, out_channel=3, filters=[28, 36, 48, 64, 80], ks=[3, 3, 5, 3, 3], blocks=[2, 2, 2, 2, 2], attn='squeeze_excitation', is_isotropic=False, isotropy=[False, False, False, True, True], pad_mode='replicate', act_mode='elu', norm_mode='bn', init_mode='orthogonal', deploy=False, fmap_size=[17, 129, 129], **kwargs)

3D feature pyramid network (FPN). This design is flexible in handling both isotropic data and anisotropic data.

Parameters

• backbone_type (str) – the block type at each U-Net stage. Default: 'resnet'

• block_type (str) – the block type in the backbone. Default: 'residual'

• in_channel (int) – number of input channels. Default: 1

• out_channel (int) – number of output channels. Default: 3

• filters (List[int]) – number of filters at each FPN stage. Default: [28, 36, 48, 64, 80]

• is_isotropic (bool) – whether the whole model is isotropic. Default: False

• isotropy (List[bool]) – specify each U-Net stage is isotropic or anisotropic. All elements will be True if is_isotropic is True. Default: [False, False, False, True, True]

• pad_mode (str) – one of 'zeros', 'reflect', 'replicate' or 'circular'. Default: 'replicate'

• act_mode (str) – one of 'relu', 'leaky_relu', 'elu', 'gelu', 'swish', 'efficient_swish' or 'none'. Default: 'relu'

• norm_mode (str) – one of 'bn', 'sync_bn' 'in' or 'gn'. Default: 'bn'

• init_mode (str) – one of 'xavier', 'kaiming', 'selu' or 'orthogonal'. Default: 'orthogonal'

• deploy (bool) – build backbone in deploy mode (exclusive for RepVGG backbone). Default: False

• feature_keys (List[str]) –

• ks (List[int]) –

• blocks (List[int]) –

• attn (str) –

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.SwinUNETR(img_size, in_channel, out_channel, depths=(2, 2, 2, 2), num_heads=(3, 6, 12, 24), feature_size=24, norm_name='instance', drop_rate=0.0, attn_drop_rate=0.0, dropout_path_rate=0.0, normalize=True, use_checkpoint=False, spatial_dims=3, downsample='merging', **kwargs)

Swin UNETR based on: “Hatamizadeh et al., Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images <https://arxiv.org/abs/2201.01266>”

Parameters

• img_size (Union[Sequence[int], int]) – dimension of input image.

• in_channels – dimension of input channels.

• out_channels – dimension of output channels.

• feature_size (int) – dimension of network feature size.

• depths (Sequence[int]) – number of layers in each stage.

• num_heads (Sequence[int]) – number of attention heads.

• norm_name (Union[Tuple, str]) – feature normalization type and arguments.

• drop_rate (float) – dropout rate.

• attn_drop_rate (float) – attention dropout rate.

• dropout_path_rate (float) – drop path rate.

• normalize (bool) – normalize output intermediate features in each stage.

• use_checkpoint (bool) – use gradient checkpointing for reduced memory usage.

• spatial_dims (int) – number of spatial dims.

• downsample – module used for downsampling, available options are “mergingv2”, “merging” and a user-specified nn.Module following the API defined in monai.networks.nets.PatchMerging. The default is currently “merging”

• in_channel (int) –

• out_channel (int) –

Return type

None

forward(x_in)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.UNETR(in_channel, out_channel, img_size, feature_size=16, hidden_size=768, mlp_dim=3072, num_heads=12, pos_embed='perceptron', norm_name='instance', conv_block=False, res_block=True, dropout_rate=0.0, **kwargs)

Based on: “Hatamizadeh et al.,<https://arxiv.org/abs/2103.10504>”.

Args:

in_channels (int): dimension of input channels. out_channels (int): dimension of output channels. img_size (Tuple[int, int, int]): dimension of input image. feature_size (int): dimension of network feature size. hidden_size (int): dimension of hidden layer. mlp_dim (int): dimension of feedforward layer. num_heads (int): number of attention heads. pos_embed (str): position embedding layer type. norm_name (Union[Tuple, str]): feature normalization type and arguments. conv_block (bool): bool argument to determine if convolutional block is used. res_block (bool): bool argument to determine if residual block is used. dropout_rate (float): faction of the input units to drop.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters

• in_channel (int) –

• out_channel (int) –

• img_size (Tuple[int, int, int]) –

• feature_size (int) –

• hidden_size (int) –

• mlp_dim (int) –

• num_heads (int) –

• pos_embed (str) –

• norm_name (Union[Tuple, str]) –

• conv_block (bool) –

• res_block (bool) –

• dropout_rate (float) –

Return type

None

forward(x_in)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.UNet2D(block_type='residual', in_channel=1, out_channel=3, filters=[32, 64, 128, 256, 512], pad_mode='replicate', act_mode='leaky_relu', norm_mode='gn', init_mode='orthogonal', pooling=False, **kwargs)

2D residual U-Net architecture.

Parameters

• block_type (str) – the block type at each U-Net stage. Default: 'residual'

• in_channel (int) – number of input channels. Default: 1

• out_channel (int) – number of output channels. Default: 3

• filters (List[int]) – number of filters at each U-Net stage. Default: [32, 64, 128, 256, 512]

• pad_mode (str) – one of 'zeros', 'reflect', 'replicate' or 'circular'. Default: 'replicate'

• act_mode (str) – one of 'relu', 'leaky_relu', 'elu', 'gelu', 'swish', 'efficient_swish' or 'none'. Default: 'leaky_relu'

• norm_mode (str) – one of 'bn', 'sync_bn' 'in' or 'gn'. Default: 'gn'

• init_mode (str) – one of 'xavier', 'kaiming', 'selu' or 'orthogonal'. Default: 'orthogonal'

• pooling (bool) – downsample by max-pooling if True else using stride. Default: False

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.UNet3D(block_type='residual', in_channel=1, out_channel=3, filters=[28, 36, 48, 64, 80], is_isotropic=False, isotropy=[False, False, False, True, True], pad_mode='replicate', act_mode='elu', norm_mode='bn', init_mode='orthogonal', pooling=False, blurpool=False, return_feats=None, **kwargs)

3D residual U-Net architecture. This design is flexible in handling both isotropic data and anisotropic data.

Parameters

• block_type (str) – the block type at each U-Net stage. Default: 'residual'

• in_channel (int) – number of input channels. Default: 1

• out_channel (int) – number of output channels. Default: 3

• filters (List[int]) – number of filters at each U-Net stage. Default: [28, 36, 48, 64, 80]

• is_isotropic (bool) – whether the whole model is isotropic. Default: False

• isotropy (List[bool]) – specify each U-Net stage is isotropic or anisotropic. All elements will be True if is_isotropic is True. Default: [False, False, False, True, True]

• pad_mode (str) – one of 'zeros', 'reflect', 'replicate' or 'circular'. Default: 'replicate'

• act_mode (str) – one of 'relu', 'leaky_relu', 'elu', 'gelu', 'swish', 'efficient_swish' or 'none'. Default: 'relu'

• norm_mode (str) – one of 'bn', 'sync_bn' 'in' or 'gn'. Default: 'bn'

• init_mode (str) – one of 'xavier', 'kaiming', 'selu' or 'orthogonal'. Default: 'orthogonal'

• pooling (bool) – downsample by max-pooling if True else using stride. Default: False

• blurpool (bool) – apply blurpool as in Zhang 2019 (https://arxiv.org/abs/1904.11486). Default: False

• return_feats (Optional[list]) –

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.UNetPlus2D(filters=[32, 64, 128, 256, 512], norm_mode='gn', **kwargs)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters

• filters (List[int]) –

• norm_mode (str) –

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class connectomics.model.arch.UNetPlus3D(filters=[28, 36, 48, 64, 80], norm_mode='bn', **kwargs)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters

• filters (List[int]) –

• norm_mode (str) –

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Loss Functions

class connectomics.model.loss.Criterion(device, target_opt=['1'], loss_opt=[['WeightedBCE']], output_act=[['none']], loss_weight=[[1.0]], loss_kwargs=None, regu_opt=None, regu_target=None, regu_weight=None, do_2d=False)

Calculating losses and regularizations given the prediction, target and weight mask.

Parameters

• device (torch.device) – model running device. GPUs are recommended for model training and inference.

• target_opt (List[str], optional) – target options. Defaults to [‘1’].

• loss_opt (List[List[str]], optional) – loss options for the specified targets. Defaults to [[‘WeightedBCE’]].

• output_act (List[List[str]], optional) – activation functions for each loss option. Defaults to [[‘none’]].

• loss_weight (List[List[float]], optional) – the scalar weight of each loss. Defaults to [[1.]].

• Optional[List[List[dict]]] (loss_kwargs) – a list of kwargs given to the loss functions. Defaults to None.

• regu_opt (Optional[List[str]], optional) – regularization options. Defaults to None.

• regu_target (Optional[List[List[int]]], optional) – indicies of predictions for applying regularization. Defaults to None.

• regu_weight (Optional[List[float]], optional) – the scalar weight of each regularization. Defaults to None.

• do_2d (bool, optional) – whether to conduct 2D training. Defaults to False.

• loss_kwargs (Optional[List[List[dict]]]) –

classmethod build_from_cfg(cfg, device)

Build a Criterion class based on the config options.

Parameters

• cfg (yacs.config.CfgNode) – YACS configuration options.

• device (torch.device) – model running device type. GPUs are recommended for model training and inference.

class connectomics.model.loss.GANLoss(gan_mode='lsgan', target_real_label=1.0, target_fake_label=0.0)

Define different GAN objectives (vanilla, lsgan, and wgangp). Based on Based on https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix

The GANLoss class abstracts away the need to create the target label tensor that has the same size as the input.

Initialize the GANLoss class.

Parameters

• gan_mode (str) – the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.

• target_real_label (bool) – label for a real image

• target_fake_label (bool) – label of a fake image

Note: Do not use sigmoid as the last layer of Discriminator. LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.

get_target_tensor(prediction, target_is_real)

Create label tensors with the same size as the input.

Parameters

• prediction (torch.Tensor) – tpyically the prediction from a discriminator

• target_is_real (bool) – if the ground truth label is for real images or fake images

Returns

A label tensor filled with ground truth label, and with the size of the input

Utility Functions

class connectomics.model.utils.ImagePool(pool_size, device, on_cpu=False)

This class implements an image buffer that stores previously generated images. Adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/util/image_pool.py

This buffer enables us to update discriminators using a history of generated images rather than the ones produced by the latest generators.

Initialize the ImagePool class

Parameters

• pool_size (int) – the size of image buffer, if pool_size=0, no buffer will be created

• device (torch.device) – model running device. GPUs are recommended for model training and inference.

• on_cpu (bool) – whether to save image buffer on cpu to reduce GPU memory usage. Defalt: False

query(images)

Return an image from the pool.

Parameters

images – the latest generated images from the generator

Returns images from the buffer.

By 50/100, the buffer will return input images. By 50/100, the buffer will return images previously stored in the buffer, and insert the current images to the buffer.

class connectomics.model.utils.SplitActivation(target_opt=['0'], output_act=None, split_only=False, do_cat=True, do_2d=False, normalize=False)

Apply different activation functions for the outpur tensor.

Parameters

• target_opt (List[str]) –

• output_act (Optional[List[str]]) –

• split_only (bool) –

• do_cat (bool) –

• do_2d (bool) –

• normalize (bool) –

connectomics.model.utils.get_activation(activation='relu')

Get the specified activation layer.

Parameters

activation (str) – one of 'relu', 'leaky_relu', 'elu', 'gelu', 'silu', 'swish', ‘efficient_swish’`` and 'none'. Default: 'relu'

Return type

torch.nn.modules.module.Module

connectomics.model.utils.get_functional_act(activation='relu')

Get the specified activation function.

Parameters

activation (str) – one of 'relu', 'tanh', 'elu', 'sigmoid', 'softmax' and 'none'. Default: 'sigmoid'

connectomics.model.utils.get_norm_2d(norm, out_channels, bn_momentum=0.1)

Get the specified normalization layer for a 2D model.

Parameters

• norm (str) – one of 'bn', 'sync_bn' 'in', 'gn' or 'none'.

• out_channels (int) – channel number.

• bn_momentum (float) – the momentum of normalization layers.

Returns

the normalization layer

Return type

nn.Module

connectomics.model.utils.get_norm_3d(norm, out_channels, bn_momentum=0.1)

Get the specified normalization layer for a 3D model.

Parameters

• norm (str) – one of 'bn', 'sync_bn' 'in', 'gn' or 'none'.

• out_channels (int) – channel number.

• bn_momentum (float) – the momentum of normalization layers.

Returns

the normalization layer

Return type

nn.Module

connectomics.model.utils.model_init(model, mode='orthogonal')

Initialization of model weights.