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.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], 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]) –

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.UNet2D(block_type='residual', in_channel=1, out_channel=3, filters=[32, 64, 128, 256, 512], pad_mode='replicate', act_mode='elu', norm_mode='bn', init_mode='orthogonal', pooling=False, **kwargs)

2D residual U-Net architecture. :param block_type: the block type at each U-Net stage. Default: 'residual' :type block_type: str :param in_channel: number of input channels. Default: 1 :type in_channel: int :param out_channel: number of output channels. Default: 3 :type out_channel: int :param filters: number of filters at each U-Net stage. Default: [28, 36, 48, 64, 80] :type filters: List[int] :param pad_mode: one of 'zeros', 'reflect', 'replicate' or 'circular'. Default: 'replicate' :type pad_mode: str :param act_mode: one of 'relu', 'leaky_relu', 'elu', 'gelu',

'swish', 'efficient_swish' or 'none'. Default: 'relu'

Parameters

• 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

• in_channel (int) –

• out_channel (int) –

• filters (List[int]) –

• pad_mode (str) –

• act_mode (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.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, **kwargs)

3D residual U-Net architecture. This design is flexible in handling both isotropic data and anisotropic data. :param block_type: the block type at each U-Net stage. Default: 'residual' :type block_type: str :param in_channel: number of input channels. Default: 1 :type in_channel: int :param out_channel: number of output channels. Default: 3 :type out_channel: int :param filters: number of filters at each U-Net stage. Default: [28, 36, 48, 64, 80] :type filters: List[int] :param is_isotropic: whether the whole model is isotropic. Default: False :type is_isotropic: bool :param isotropy: 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]

Parameters

• 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

• in_channel (int) –

• out_channel (int) –

• filters (List[int]) –

• is_isotropic (bool) –

• isotropy (List[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.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]], 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.]].

• 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.

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.

Utility Functions

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_functional_act(activation='relu')

Get the specified activation function.

Parameters

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