from fastai.vision.all import test_eqNetworks
Neural networks
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.
Args: spatial_dims (int, optional): The number of spatial dimensions. Default is 2. in_channels (int, optional): The number of input channels. Default is 1 (grayscale image). out_channels (int, optional): The number of output channels. Default is 1 (denoised image). num_of_layers (int, optional): The number of convolutional layers in the network. Default is 9. features (int, optional): The number of feature maps in the first layer and subsequent layers. Default is 64. kernel_size (int or tuple, optional): The size of the convolution kernel. Default is 3.*
x = torch_randn(16, 1, 32, 64)
tst = DnCNN(2,1)
test_eq(tst(x).shape, x.shape)DeepLab v3+
Config
interpolate
interpolate (x:torch.Tensor, size:Union[List[int],Tuple[int,...]], dims:int)
get_padding
get_padding (kernel_size:int, dilation:int)
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
Block
Block (config:__main__.DeeplabConfig, inplanes:int, planes:int, reps:int, stride:int=1, dilation:int=1, start_with_relu:bool=True, grow_first:bool=True, is_last:bool=False)
*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*
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*
Aligned Xception
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*
ASPP
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*
DeepLab V3
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_graddef 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.
# Args:
# dim (int): The dimension of the input tensor. This is typically the feature size.
# d_state (int, optional): Expansion factor for the state in the Mamba block. Default is 16.
# d_conv (int, optional): Width of the local convolution in the Mamba block. Default is 4.
# expand (int, optional): Factor by which to expand the dimensions in the Mamba block. Default is 2.
# Attributes:
# dim (int): The dimension of the input tensor.
# norm (nn.LayerNorm): Layer normalization applied to the input tensor before processing.
# mamba (Mamba): The core Mamba block which performs spatial-spectral transformations.
# Methods:
# forward(x):
# Defines the computation performed at every call.
# Args:
# x (torch.Tensor): Input tensor of shape [batch_size, dim, height, width].
# Returns:
# torch.Tensor: Output tensor after applying Mamba block and normalization.
# """
# def __init__(self, dim, d_state=16, d_conv=4, expand=2):
# super().__init__()
# self.dim = dim
# self.norm = nn.LayerNorm(dim)
# self.mamba = Mamba(
# 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):
# """
# Forward pass of the MambaLayer. Applies layer normalization and optionally converts input precision.
# Args:
# x (torch.Tensor): Input tensor of shape [batch_size, dim, height, width].
# 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# #| export
# class UMamba(DynUNet):
# """
# 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.
# Methods:
# get_bottleneck():
# Constructs and returns the bottleneck part of the network, which includes a convolution block followed by a MambaLayer.
# Returns:
# nn.Sequential: A PyTorch sequential container holding the convolution block and the MambaLayer for the bottleneck.
# """
# def get_bottleneck(self):
# """
# Constructs the bottleneck part of the network.
# The bottleneck consists of a convolution block followed by a MambaLayer. Both components are added to a sequential container.
# Returns:
# nn.Sequential: A PyTorch sequential container with the convolution block and MambaLayer for the bottleneck.
# """
# mamba_bottleneck = []
# # Add a convolution block before the MambaLayer in the bottleneck
# mamba_bottleneck.append(
# self.conv_block(
# self.spatial_dims, # Spatial dimensions of the input data
# self.filters[-2], # Number of filters for the previous layer
# self.filters[-1], # Number of filters for this layer (output)
# self.kernel_size[-1], # Kernel size for the convolution
# self.strides[-1], # Stride for the convolution
# self.norm_name, # Normalization method name
# self.act_name, # Activation function name
# dropout=self.dropout # Dropout probability
# )
# )
# # Add the MambaLayer to the bottleneck
# mamba_bottleneck.append(
# MambaLayer(dim = self.filters[-1]) # Initialize the MambaLayer with the current filter size as dimension
# )
# return nn.Sequential(*mamba_bottleneck) # Return the sequential container holding both componentsExample
# 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)# 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)