Networks

Neural networks

UNet

---

create_custom_unet

 create_custom_unet (resnet_version, n_in=1, n_out=1, img_size=(128, 128),
                     pretrained=True, cut=4)

Create a U-Net model with a ResNet backbone.

Returns: - U-Net model with the specified ResNet backbone.

	Type	Default	Details
resnet_version			Choose a ResNet model between: ‘resnet18’, ‘resnet34’, ‘resnet50’, ‘resnet101’, and ‘resnet152’.
n_in	int	1	Number of input channels, default is 1 (e.g., grayscale).
n_out	int	1	Number of output channels.
img_size	tuple	(128, 128)	Tuple for the input image size, default is (128, 128).
pretrained	bool	True	If True, use a pretrained ResNet backbone.
cut	int	4	The cut point for the ResNet model, default is 4.

Denoising CNN

---

DnCNN

 DnCNN (spatial_dims=2, in_channels=1, out_channels=1, num_of_layers=9,
        features=64, kernel_size=3)

A Deep Neural Network for Image Denoising (DnCNN) model.

	Type	Default	Details
spatial_dims	int	2	Number of spatial dimensions
in_channels	int	1	Number of input channels
out_channels	int	1	Number of output channels
num_of_layers	int	9	Number of convolutional layers
features	int	64	Number of feature maps
kernel_size	int	3	Size of the convolution kernel

x = torch_randn(16, 1, 32, 64)

tst = DnCNN(2,1)
test_eq(tst(x).shape, x.shape)

DeepLab v3+

DeepLabv3 is a semantic segmentation architecture that handles the problem of segmenting objects at multiple scales. It uses the Atroys Spatial Pyramid Pooling module, and introduces various updates with respect to other versions.

Config

The DeeplabConfig class has been created to centralize all settings and hyperparameters in one place. It uses two main functions: Get_padding and interpolate. Get_padding is a function that calculates the amount of padding needed for a convolution operation to get the desired output size. Interpolate is a function that allows the resizing of a tensor using interpolation.

---

interpolate

 interpolate (x:torch.Tensor, size:Union[List[int],Tuple[int,...]],
              dims:int)

	Type	Details
x	Tensor	Input tensor
size	Union	Size of the output tensor
dims	int	Number of spatial dimensions
Returns	Tensor	Output tensor

---

get_padding

 get_padding (kernel_size:int, dilation:int)

	Type	Details
kernel_size	int	Size of the convolution kernel
dilation	int	Dilation rate
Returns	int	Padding size

---

DeeplabConfig

 DeeplabConfig (dimensions:int, in_channels:int, out_channels:int,
                backbone:str='xception', pretrained:bool=False,
                middle_flow_blocks:int=16,
                aspp_dilations:List[int]=<factory>,
                entry_block3_stride:int=2, middle_block_dilation:int=1,
                exit_block_dilations:Tuple[int,int]=(1, 2))

Blocks

SeparableConv is a class that carries out a type of convolution operation that splits the traditional convolution into two parts: depthwise convolution (a convolution filter for each channel independently), and pointwise convolution (combines the outputs of the depthwise convolution).

---

SeparableConv

 SeparableConv (config:__main__.DeeplabConfig, inplanes:int, planes:int,
                kernel_size:int=3, stride:int=1, dilation:int=1,
                bias:bool=False, norm:Optional[str]=None)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.

ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

	Type	Default	Details
config	DeeplabConfig		Configuration for the Deeplab model
inplanes	int		Number of input channels
planes	int		Number of output channels
kernel_size	int	3	Size of the convolution kernel
stride	int	1	Stride for the convolution
dilation	int	1	Dilation rate for the convolution
bias	bool	False	If True, add a bias term
norm	Optional	None	Type of normalization layer

Block is a class that combines multiple separable convolutional layers and residual connections.

class Block(nn.Module):
    def __init__(self, config: DeeplabConfig, # Configuration for the Deeplab model
                 inplanes: int, # Number of input channels
                 planes: int, # Number of output channels
                 reps: int, # Number of convolutional layers
                 stride: int = 1,  # Stride for the convolution
                 dilation: int = 1, # Dilation rate for the convolution
                 start_with_relu: bool = True, # If True, start with a ReLU activation
                 grow_first: bool = True, # If True, increase the number of channels in the first convolution
                 is_last: bool = False, # If True, add a convolution layer at the end
                 ):
        super().__init__()
        if planes != inplanes or stride != 1:
            self.skip = Convolution(config.dimensions, inplanes, planes, kernel_size=1, bias=False, 
                                    strides=stride, norm=Norm.BATCH)
        else:
            self.skip = None

        self.relu = nn.ReLU(inplace=True)
        rep = []

        filters = inplanes
        if grow_first:
            rep.append(self.relu)
            rep.append(SeparableConv(config, inplanes, planes, 3, stride=1, dilation=dilation, norm=Norm.BATCH))
            filters = planes

        for _ in range(reps - 1):
            rep.append(self.relu)
            rep.append(SeparableConv(config, filters, filters, 3, stride=1, dilation=dilation, norm=Norm.BATCH))

        if not grow_first:
            rep.append(self.relu)
            rep.append(SeparableConv(config, inplanes, planes, 3, stride=1, dilation=dilation, norm=Norm.BATCH))

        if not start_with_relu:
            rep = rep[1:]

        if stride != 1:
            rep.append(SeparableConv(config, planes, planes, 3, stride=2))

        if stride == 1 and is_last:
            rep.append(SeparableConv(config, planes, planes, 3, stride=1))

        self.rep = nn.Sequential(*rep)

    def forward(self, inp: torch_Tensor) -> torch_Tensor:
        x = self.rep(inp)
        if self.skip is not None:
            skip = self.skip(inp)
        else:
            skip = inp
        x += skip
        return x

Aligned Xception

The Xception class defines the Xception Neural Network used in DeepLab.

---

Xception

 Xception (config:__main__.DeeplabConfig)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.

ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

	Type	Details
config	DeeplabConfig	Configuration for the Deeplab model

ASPP

The class ASPP_module to compute the Atroys Spatial Pyramid Pooling method and create the convolution that uses it.

---

ASPP_module

 ASPP_module (config:__main__.DeeplabConfig, inplanes:int, planes:int,
              dilation:int)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.

ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

	Type	Details
config	DeeplabConfig	Configuration for the Deeplab model
inplanes	int	Number of input channels
planes	int	Number of output channels
dilation	int	Dilation rate for the convolution

DeepLab V3

The Deeplab class combines the different modules to make the DeepLab V3 architecture.

---

Deeplab

 Deeplab (config:__main__.DeeplabConfig)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Example

# Load a pre-trained ResNet backbone
resnet_backbone = ResNetFeatures('resnet10', pretrained=False, in_channels=1, spatial_dims=3)

# Forward pass through the backbone to get the output before the final classifier
dummy_input = torch_randn(1, 1, 64, 224, 224)  # Example input size; adjust based on your needs
output = resnet_backbone(dummy_input)

# The shape of 'output' will give you the number of channels at this stage in the backbone
print("Output channels:", output[-1].shape[1])

Output channels: 512

# For 2D images
config_2d = DeeplabConfig(
    dimensions=2,
    in_channels=3,  # For RGB images
    out_channels=4,
    backbone="xception",  # or whatever backbone you're using
    aspp_dilations=[1, 6, 12, 18]
)
model_2d = Deeplab(config_2d)

# For 3D images
config_3d = DeeplabConfig(
    dimensions=3,
    in_channels=1,  # For single-channel 3D medical images
    out_channels=4,
    middle_flow_blocks=16,
    aspp_dilations=[1, 6, 12, 18]
)
model_3d = Deeplab(config_3d)

from torch import no_grad as torch_no_grad

def test_deeplab(config, input_shape, expected_output_shape):
    set_determinism(0)  # For reproducibility
    
    model = Deeplab(config)
    model.eval()  # Set the model to evaluation mode
    
    # Generate random input tensor
    x = torch_randn(*input_shape)
    
    # Forward pass
    with torch_no_grad():
        output = model(x)
    
    # Check output shape
    assert output.shape == expected_output_shape, f"Expected shape {expected_output_shape}, but got {output.shape}"
    
    print(f"Test passed for {config.dimensions}D model with backbone {config.backbone}")
    print(f"Input shape: {input_shape}")
    print(f"Output shape: {output.shape}")
    print("---")

# Test 2D model
config_2d = DeeplabConfig(
    dimensions=2,
    in_channels=3,
    out_channels=4,
    backbone="xception",
    aspp_dilations=[1, 6, 12, 18]
)
test_deeplab(config_2d, (1, 3, 64, 64), (1, 4, 64, 64))

# Test 2D model with ResNet50 backbone
config_2d_resnet = DeeplabConfig(
    dimensions=2,
    in_channels=3,
    out_channels=4,
    backbone="resnet50",
    aspp_dilations=[1, 6, 12, 18]
)
test_deeplab(config_2d_resnet, (1, 3, 64, 64), (1, 4, 64, 64))

# Test 3D model
config_3d = DeeplabConfig(
    dimensions=3,
    in_channels=1,
    out_channels=4,
    backbone="xception",
    aspp_dilations=[1, 6, 12, 18]
)
test_deeplab(config_3d, (1, 1, 64, 64, 64), (1, 4, 64, 64, 64))

# Test 3D model with ResNet10 backbone
config_3d_resnet = DeeplabConfig(
    dimensions=3,
    in_channels=1,
    out_channels=4,
    backbone="resnet10",
    aspp_dilations=[1, 6, 12, 18]
)
test_deeplab(config_3d_resnet, (1, 1, 64, 64, 64), (1, 4, 64, 64, 64))

print("All tests passed successfully!")

Test passed for 2D model with backbone xception
Input shape: (1, 3, 64, 64)
Output shape: torch.Size([1, 4, 64, 64])
---
Test passed for 2D model with backbone resnet50
Input shape: (1, 3, 64, 64)
Output shape: torch.Size([1, 4, 64, 64])
---
Test passed for 3D model with backbone xception
Input shape: (1, 1, 64, 64, 64)
Output shape: torch.Size([1, 4, 64, 64, 64])
---
Test passed for 3D model with backbone resnet10
Input shape: (1, 1, 64, 64, 64)
Output shape: torch.Size([1, 4, 64, 64, 64])
---
All tests passed successfully!

UMamba

---

MambaLayer

 MambaLayer (dim, d_state=16, d_conv=4, expand=2)

A custom neural network layer that incorporates the Mamba block from the Mamba model, along with layer normalization and optional mixed precision handling.

	Type	Default	Details
dim			Dimension of the input tensor
d_state	int	16	Expansion factor for the state in the Mamba block
d_conv	int	4	Width of the local convolution in the Mamba block
expand	int	2	Factor by which to expand the dimensions in the Mamba block

---

UMamba

 UMamba (spatial_dims:int, in_channels:int, out_channels:int,
         kernel_size:Sequence[Union[Sequence[int],int]],
         strides:Sequence[Union[Sequence[int],int]],
         upsample_kernel_size:Sequence[Union[Sequence[int],int]],
         filters:Optional[Sequence[int]]=None,
         dropout:Union[Tuple,str,float,NoneType]=None,
         norm_name:Union[Tuple,str]=('INSTANCE', {'affine': True}),
         act_name:Union[Tuple,str]=('leakyrelu', {'inplace': True,
         'negative_slope': 0.01}), deep_supervision:bool=False,
         deep_supr_num:int=1, res_block:bool=False, trans_bias:bool=False)

A custom subclass of DynUNet that integrates the Mamba layer into the model’s bottleneck.

This class inherits from DynUNet and adds a specific bottleneck structure containing a convolution block followed by a MambaLayer.

Example

x = torch_randn(16, 1, 32, 64)

tst = DynUNet(2,1,1,[3,3,3],[1,1,1],[1,1])
print(tst(x).shape)
test_eq(tst(x).shape, x.shape)

torch.Size([16, 1, 32, 64])

x = torch_randn(16, 1, 32, 64).cuda()

tst = UMamba(2,1,1,[3,3,3],[1,1,1],[1,1]).cuda()
print(tst(x).shape)
test_eq(tst(x).shape, x.shape)

torch.Size([16, 1, 32, 64])

---