Networks

Neural networks

UNet

---

source

create_custom_unet

def create_custom_unet(
    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.
):

Create a U-Net model with a ResNet backbone.

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

Denoising CNN

---

source

DnCNN

def DnCNN(
    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
):

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

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.

---

source

interpolate

def interpolate(
    x:Tensor, # Input tensor
    size:Union, # Size of the output tensor
    dims:int, # Number of spatial dimensions
)->Tensor: # Output tensor

---

source

get_padding

def get_padding(
    kernel_size:int, # Size of the convolution kernel
    dilation:int, # Dilation rate
)->int: # Padding size

---

source

DeeplabConfig

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

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

---

source

SeparableConv

def SeparableConv(
    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
):

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested 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) -> None:
        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 also have their parameters converted 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

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.

---

source

Xception

def Xception(
    config:DeeplabConfig, # Configuration for the Deeplab model
):

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested 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) -> None:
        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 also have their parameters converted 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

ASPP

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

---

source

ASPP_module

def ASPP_module(
    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
):

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested 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) -> None:
        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 also have their parameters converted 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

DeepLab V3

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

---

source

Deeplab

def Deeplab(
    config:DeeplabConfig
):

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested 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) -> None:
        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 also have their parameters converted 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

# #| export
# class MambaLayer(nn.Module):
#     """
#     A custom neural network layer that incorporates the Mamba block from the Mamba model, 
#     along with layer normalization and optional mixed precision handling.

#     """
    
#     def __init__(self, dim,  # Dimension of the input tensor
#                  d_state=16, # Expansion factor for the state in the Mamba block
#                  d_conv=4,   # Width of the local convolution in the Mamba block
#                  expand=2,   # Factor by which to expand the dimensions in the Mamba block
#                  ):
#         super().__init__()
#         self.dim = dim                      # Dimension of the input tensor
#         self.norm = nn.LayerNorm(dim)       # Layer normalization
#         self.mamba = Mamba(                 # Mamba block
#             d_model=dim,  # Model dimension d_model
#             d_state=d_state,  # SSM state expansion factor
#             d_conv=d_conv,  # Local convolution width
#             expand=expand  # Block expansion factor
#         )
    
#     @autocast(enabled=False)
#     def forward(self, x, # Input tensor of shape [batch_size, dim, height, width].
#                 ):
#         """
#         Forward pass of the MambaLayer. Applies layer normalization and optionally converts input precision.
   
#         Returns:
#             torch.Tensor: Output tensor after applying Mamba block and normalization.
#         """
#         if x.dtype == torch_float16:
#             x = x.type(torch_float32)  # Convert input to float32 for mixed precision handling
        
#         B, C = x.shape[:2]
#         assert C == self.dim  # Ensure the feature size matches the dimension of the layer
        
#         n_tokens = x.shape[2:].numel()
#         img_dims = x.shape[2:]
#         x_flat = x.reshape(B, C, n_tokens).transpose(-1, -2)  # Flatten and transpose for Mamba input
        
#         x_norm = self.norm(x_flat)  # Apply layer normalization
        
#         x_mamba = self.mamba(x_norm)  # Pass through the Mamba block
        
#         out = x_mamba.transpose(-1, -2).reshape(B, C, *img_dims)  # Reshape and transpose back to original dimensions
        
#         return out

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).to(device)

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

---