Commit 12ce71d4 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

shallow dense block + competitive dense blocks architectures

parent bd826338
......@@ -20,6 +20,379 @@ import torch.nn as nn
import utils.modules as modules
class BrainMapperCompResUNet3D(nn.Module):
"""Architecture class for Competitive Residual DenseBlock BrainMapper 3D U-net.
This class contains the pytorch implementation of the U-net architecture underpinning the BrainMapper project.
Args:
parameters (dict): Contains information relevant parameters
parameters = {
'kernel_heigth': 5
'kernel_width': 5
'kernel_depth': 5
'kernel_classification': 1
'input_channels': 1
'output_channels': 64
'convolution_stride': 1
'dropout': 0.2
'pool_kernel_size': 2
'pool_stride': 2
'up_mode': 'upconv'
'number_of_classes': 1
}
Returns:
probability_map (torch.tensor): Output forward passed tensor through the U-net block
"""
def __init__(self, parameters):
super(BrainMapperCompResUNet3D, self).__init__()
original_input_channels = parameters['input_channels']
original_output_channels = parameters['output_channels']
self.encoderBlock1 = modules.InCompDensEncoderBlock3D(parameters)
parameters['input_channels'] = parameters['output_channels']
self.encoderBlock2 = modules.CompDensEncoderBlock3D(parameters)
self.encoderBlock3 = modules.CompDensEncoderBlock3D(parameters)
self.encoderBlock4 = modules.CompDensEncoderBlock3D(parameters)
self.bottleneck = modules.CompDensBlock3D(parameters)
self.decoderBlock1 = modules.CompDensDecoderBlock3D(parameters)
self.decoderBlock2 = modules.CompDensDecoderBlock3D(parameters)
self.decoderBlock3 = modules.CompDensDecoderBlock3D(parameters)
self.decoderBlock4 = modules.CompDensDecoderBlock3D(parameters)
self.classifier = modules.DensClassifierBlock3D(parameters)
parameters['input_channels'] = original_input_channels
parameters['output_channels'] = original_output_channels
def forward(self, X):
"""Forward pass for 3D U-net
Function computing the forward pass through the 3D U-Net
The input to the function is the dMRI map
Args:
X (torch.tensor): Input dMRI map, shape = (N x C x D x H x W)
Returns:
probability_map (torch.tensor): Output forward passed tensor through the U-net block
"""
Y_encoder_1, Y_np1, _ = self.encoderBlock1.forward(X)
Y_encoder_2, Y_np2, _ = self.encoderBlock2.forward(
Y_encoder_1)
del Y_encoder_1
Y_encoder_3, Y_np3, _ = self.encoderBlock3.forward(
Y_encoder_2)
del Y_encoder_2
Y_encoder_4, Y_np4, _ = self.encoderBlock4.forward(
Y_encoder_3)
del Y_encoder_3
Y_bottleNeck = self.bottleneck.forward(Y_encoder_4)
del Y_encoder_4
Y_decoder_1 = self.decoderBlock1.forward(
Y_bottleNeck, Y_np4)
del Y_bottleNeck, Y_np4
Y_decoder_2 = self.decoderBlock2.forward(
Y_decoder_1, Y_np3)
del Y_decoder_1, Y_np3
Y_decoder_3 = self.decoderBlock3.forward(
Y_decoder_2, Y_np2)
del Y_decoder_2, Y_np2
Y_decoder_4 = self.decoderBlock4.forward(
Y_decoder_3, Y_np1)
del Y_decoder_3, Y_np1
probability_map = self.classifier.forward(Y_decoder_4)
del Y_decoder_4
return probability_map
def save(self, path):
"""Model Saver
Function saving the model with all its parameters to a given path.
The path must end with a *.model argument.
Args:
path (str): Path string
"""
print("Saving Model... {}".format(path))
torch.save(self, path)
@property
def test_if_cuda(self):
"""Cuda Test
This function tests if the model parameters are allocated to a CUDA enabled GPU.
Returns:
bool: Flag indicating True if the tensor is stored on the GPU and Flase otherwhise
"""
return next(self.parameters()).is_cuda
def predict(self, X, device=0):
"""Post-training Output Prediction
This function predicts the output of the of the U-net post-training
Args:
X (torch.tensor): input dMRI volume
device (int/str): Device type used for training (int - GPU id, str- CPU)
Returns:
prediction (ndarray): predicted output after training
"""
self.eval() # PyToch module setting network to evaluation mode
if type(X) is np.ndarray:
X = torch.tensor(X, requires_grad=False).type(torch.FloatTensor)
elif type(X) is torch.Tensor and not X.is_cuda:
X = X.type(torch.FloatTensor).cuda(device, non_blocking=True)
# .cuda() call transfers the densor from the CPU to the GPU if that is the case.
# Non-blocking argument lets the caller bypas synchronization when necessary
with torch.no_grad(): # Causes operations to have no gradients
output = self.forward(X)
_, idx = torch.max(output, 1)
# We retrieve the tensor held by idx (.data), and map it to a cpu as an ndarray
idx = idx.data.cpu().numpy()
prediction = np.squeeze(idx)
del X, output, idx
return prediction
def reset_parameters(self):
"""Parameter Initialization
This function (re)initializes the parameters of the defined network.
This function is a wrapper for the reset_parameters() function defined for each module.
More information can be found here: https://discuss.pytorch.org/t/what-is-the-default-initialization-of-a-conv2d-layer-and-linear-layer/16055 + https://discuss.pytorch.org/t/how-to-reset-model-weights-to-effectively-implement-crossvalidation/53859
An alternative (re)initialization method is described here: https://discuss.pytorch.org/t/how-to-reset-variables-values-in-nn-modules/32639
"""
print("Initializing network parameters...")
for _, module in self.named_children():
for _, submodule in module.named_children():
for _, subsubmodule in submodule.named_children():
if isinstance(subsubmodule, (torch.nn.PReLU, torch.nn.Dropout3d, torch.nn.MaxPool3d)) == False:
subsubmodule.reset_parameters()
print("Initialized network parameters!")
class BrainMapperResUNet3Dshallow(nn.Module):
"""Architecture class for Residual DenseBlock BrainMapper 3D U-net.
This class contains the pytorch implementation of the U-net architecture underpinning the BrainMapper project.
Args:
parameters (dict): Contains information relevant parameters
parameters = {
'kernel_heigth': 5
'kernel_width': 5
'kernel_depth': 5
'kernel_classification': 1
'input_channels': 1
'output_channels': 64
'convolution_stride': 1
'dropout': 0.2
'pool_kernel_size': 2
'pool_stride': 2
'up_mode': 'upconv'
'number_of_classes': 1
}
Returns:
probability_map (torch.tensor): Output forward passed tensor through the U-net block
"""
def __init__(self, parameters):
super(BrainMapperResUNet3Dshallow, self).__init__()
original_input_channels = parameters['input_channels']
original_output_channels = parameters['output_channels']
self.encoderBlock1 = modules.DensEncoderBlock3D(parameters)
parameters['input_channels'] = parameters['output_channels']
self.encoderBlock2 = modules.DensEncoderBlock3D(parameters)
self.encoderBlock3 = modules.DensEncoderBlock3D(parameters)
self.bottleneck = modules.DensBlock3D(parameters)
parameters['input_channels'] = parameters['output_channels'] * 2
self.decoderBlock1 = modules.DensDecoderBlock3D(parameters)
self.decoderBlock2 = modules.DensDecoderBlock3D(parameters)
self.decoderBlock3 = modules.DensDecoderBlock3D(parameters)
parameters['input_channels'] = parameters['output_channels']
self.classifier = modules.DensClassifierBlock3D(parameters)
parameters['input_channels'] = original_input_channels
parameters['output_channels'] = original_output_channels
def forward(self, X):
"""Forward pass for 3D U-net
Function computing the forward pass through the 3D U-Net
The input to the function is the dMRI map
Args:
X (torch.tensor): Input dMRI map, shape = (N x C x D x H x W)
Returns:
probability_map (torch.tensor): Output forward passed tensor through the U-net block
"""
Y_encoder_1, Y_np1, _ = self.encoderBlock1.forward(X)
Y_encoder_2, Y_np2, _ = self.encoderBlock2.forward(
Y_encoder_1)
del Y_encoder_1
Y_encoder_3, Y_np3, _ = self.encoderBlock3.forward(
Y_encoder_2)
del Y_encoder_2
Y_bottleNeck = self.bottleneck.forward(Y_encoder_3)
del Y_encoder_3
Y_decoder_1 = self.decoderBlock1.forward(
Y_bottleNeck, Y_np3)
del Y_bottleNeck, Y_np3
Y_decoder_2 = self.decoderBlock2.forward(
Y_decoder_1, Y_np2)
del Y_decoder_1, Y_np2
Y_decoder_3 = self.decoderBlock3.forward(
Y_decoder_2, Y_np1)
del Y_decoder_2, Y_np1
probability_map = self.classifier.forward(Y_decoder_3)
del Y_decoder_3
return probability_map
def save(self, path):
"""Model Saver
Function saving the model with all its parameters to a given path.
The path must end with a *.model argument.
Args:
path (str): Path string
"""
print("Saving Model... {}".format(path))
torch.save(self, path)
@property
def test_if_cuda(self):
"""Cuda Test
This function tests if the model parameters are allocated to a CUDA enabled GPU.
Returns:
bool: Flag indicating True if the tensor is stored on the GPU and Flase otherwhise
"""
return next(self.parameters()).is_cuda
def predict(self, X, device=0):
"""Post-training Output Prediction
This function predicts the output of the of the U-net post-training
Args:
X (torch.tensor): input dMRI volume
device (int/str): Device type used for training (int - GPU id, str- CPU)
Returns:
prediction (ndarray): predicted output after training
"""
self.eval() # PyToch module setting network to evaluation mode
if type(X) is np.ndarray:
X = torch.tensor(X, requires_grad=False).type(torch.FloatTensor)
elif type(X) is torch.Tensor and not X.is_cuda:
X = X.type(torch.FloatTensor).cuda(device, non_blocking=True)
# .cuda() call transfers the densor from the CPU to the GPU if that is the case.
# Non-blocking argument lets the caller bypas synchronization when necessary
with torch.no_grad(): # Causes operations to have no gradients
output = self.forward(X)
_, idx = torch.max(output, 1)
# We retrieve the tensor held by idx (.data), and map it to a cpu as an ndarray
idx = idx.data.cpu().numpy()
prediction = np.squeeze(idx)
del X, output, idx
return prediction
def reset_parameters(self):
"""Parameter Initialization
This function (re)initializes the parameters of the defined network.
This function is a wrapper for the reset_parameters() function defined for each module.
More information can be found here: https://discuss.pytorch.org/t/what-is-the-default-initialization-of-a-conv2d-layer-and-linear-layer/16055 + https://discuss.pytorch.org/t/how-to-reset-model-weights-to-effectively-implement-crossvalidation/53859
An alternative (re)initialization method is described here: https://discuss.pytorch.org/t/how-to-reset-variables-values-in-nn-modules/32639
"""
print("Initializing network parameters...")
for _, module in self.named_children():
for _, submodule in module.named_children():
for _, subsubmodule in submodule.named_children():
if isinstance(subsubmodule, (torch.nn.PReLU, torch.nn.Dropout3d, torch.nn.MaxPool3d)) == False:
subsubmodule.reset_parameters()
print("Initialized network parameters!")
class BrainMapperResUNet3D(nn.Module):
"""Architecture class for Residual DenseBlock BrainMapper 3D U-net.
......
......@@ -40,7 +40,7 @@ import torch.utils.data as data
import numpy as np
from solver import Solver
from BrainMapperUNet import BrainMapperUNet3D, BrainMapperResUNet3D
from BrainMapperUNet import BrainMapperUNet3D, BrainMapperResUNet3D, BrainMapperResUNet3Dshallow, BrainMapperCompResUNet3D
from utils.data_utils import get_datasets, data_test_train_validation_split, update_shuffling_flag, create_folder
import utils.data_evaluation_utils as evaluations
from utils.data_logging_utils import LogWriter
......@@ -150,7 +150,9 @@ def train(data_parameters, training_parameters, network_parameters, misc_paramet
training_parameters['pre_trained_path'])
else:
# BrainMapperModel = BrainMapperUNet3D(network_parameters)
BrainMapperModel = BrainMapperResUNet3D(network_parameters)
# BrainMapperModel = BrainMapperResUNet3D(network_parameters)
# BrainMapperModel = BrainMapperResUNet3Dshallow(network_parameters)
BrainMapperModel = BrainMapperCompResUNet3D(network_parameters)
BrainMapperModel.reset_parameters()
......
......@@ -19,6 +19,516 @@ import torch.nn.functional as F
# TODO: Currently, it appears that we are using constant size filters. We will need to adjust this in the network architecture, to allow it to encode/decode information!
# CompetitiveResBlock 3D UNet:
class CompDensBlock3D(nn.Module):
"""Parent class for a 3D convolutional competitive residual block.
This class represents a generic parent class for a convolutional residual 3D encoder or decoder block.
The class represents a subclass/child class of nn.Module, inheriting its functionality.
Args:
parameters (dict): Contains information on kernel size, number of channels, number of filters, and if convolution is strided.
parameters = {
'kernel_heigth': 5
'kernel_width': 5
'kernel_depth' : 5
'input_channels': 64
'output_channels': 64
'convolution_stride': 1
'dropout': 0.2
}
Returns:
torch.tensor: Output forward passed tensor
"""
def __init__(self, parameters):
super(CompDensBlock3D, self).__init__()
# We first calculate the amount of zero padding required (http://cs231n.github.io/convolutional-networks/)
padding_heigth = int((parameters['kernel_heigth'] - 1) / 2)
padding_width = int((parameters['kernel_width'] - 1) / 2)
padding_depth = int((parameters['kernel_depth'] - 1) / 2)
# convolutional_layer2_input = int(parameters['input_channels'] + parameters['output_channels'])
# convolutional_layer3_input = int(convolutional_layer2_input + parameters['output_channels'])
self.convolutional_layer1 = nn.Sequential(
nn.PReLU(),
nn.Conv3d(
in_channels=parameters['input_channels'],
out_channels=parameters['output_channels'],
kernel_size=(parameters['kernel_depth'],
parameters['kernel_heigth'],
parameters['kernel_width']),
stride=parameters['convolution_stride'],
padding=(padding_depth, padding_heigth, padding_width)
),
nn.InstanceNorm3d(num_features=parameters['output_channels'])
)
self.convolutional_layer2 = nn.Sequential(
nn.PReLU(),
nn.Conv3d(
in_channels=parameters['output_channels'],
out_channels=parameters['output_channels'],
kernel_size=(parameters['kernel_depth'],
parameters['kernel_heigth'],
parameters['kernel_width']),
stride=parameters['convolution_stride'],
padding=(padding_depth, padding_heigth, padding_width)
),
nn.InstanceNorm3d(num_features=parameters['output_channels'])
)
self.convolutional_layer3 = nn.Sequential(
nn.PReLU(),
nn.Conv3d(
in_channels=parameters['output_channels'],
out_channels=parameters['output_channels'],
kernel_size=(parameters['kernel_classification'],
parameters['kernel_classification'],
parameters['kernel_classification']),
stride=parameters['convolution_stride'],
padding=(0, 0, 0)
),
nn.InstanceNorm3d(num_features=parameters['output_channels'])
)
# Other activation functions which might be interesting to test:
# More reading: https://arxiv.org/abs/1706.02515 ; https://mlfromscratch.com/activation-functions-explained/#/
# self.activation = nn.SELU()
# self.activation = nn.ELU()
# self.activation = nn.ReLU()
# Instance normalisation is used to the the small batch size, and as it has shown promise during the experiments with the simple network.
if parameters['dropout'] > 0:
self.dropout_needed = True
self.dropout = nn.Dropout3d(parameters['dropout'])
else:
self.dropout_needed = False
def forward(self, X):
"""Forward pass
Function computing the forward pass through the convolutional layer.
The input to the function is a torch tensor of shape N (batch size) x C (number of channels) x D (input depth) x H (input heigth) x W (input width)
Args:
X (torch.tensor): Input tensor, shape = (N x C x D x H x W)
Returns:
torch.tensor: Output forward passed tensor
"""
feature_map1 = self.convolutional_layer1(X)
feature_map2 = torch.max(X, feature_map1)
del X, feature_map1
feature_map3 = self.convolutional_layer2(feature_map2)
feature_map4 = torch.max(feature_map2, feature_map3)
del feature_map2, feature_map3
feature_map5 = self.convolutional_layer3(feature_map4)
del feature_map4
return feature_map5
class InCompDensBlock3D(nn.Module):
"""Parent class for the input 3D convolutional competitive residual block.
This class represents a generic parent class for a convolutional residual 3D encoder or decoder block.
The class represents a subclass/child class of nn.Module, inheriting its functionality.
Args:
parameters (dict): Contains information on kernel size, number of channels, number of filters, and if convolution is strided.
parameters = {
'kernel_heigth': 5
'kernel_width': 5
'kernel_depth' : 5
'input_channels': 64
'output_channels': 64
'convolution_stride': 1
'dropout': 0.2
}
Returns:
torch.tensor: Output forward passed tensor
"""
def __init__(self, parameters):
super(InCompDensBlock3D, self).__init__()
# We first calculate the amount of zero padding required (http://cs231n.github.io/convolutional-networks/)
padding_heigth = int((parameters['kernel_heigth'] - 1) / 2)
padding_width = int((parameters['kernel_width'] - 1) / 2)
padding_depth = int((parameters['kernel_depth'] - 1) / 2)
convolutional_layer2_input = int(parameters['input_channels'] + parameters['output_channels'])
convolutional_layer3_input = int(convolutional_layer2_input + parameters['output_channels'])
self.convolutional_layer1 = nn.Sequential(
nn.InstanceNorm3d(num_features=parameters['input_channels']),
nn.Conv3d(
in_channels=parameters['input_channels'],
out_channels=parameters['output_channels'],
kernel_size=(parameters['kernel_depth'],
parameters['kernel_heigth'],
parameters['kernel_width']),
stride=parameters['convolution_stride'],
padding=(padding_depth, padding_heigth, padding_width)
),
nn.InstanceNorm3d(num_features=parameters['output_channels'])
)
self.convolutional_layer2 = nn.Sequential(
nn.PReLU(),
nn.Conv3d(
in_channels=parameters['output_channels'],
out_channels=parameters['output_channels'],
kernel_size=(parameters['kernel_depth'],
parameters['kernel_heigth'],
parameters['kernel_width']),
stride=parameters['convolution_stride'],
padding=(padding_depth, padding_heigth, padding_width)
),
nn.InstanceNorm3d(num_features=parameters['output_channels'])
)
self.convolutional_layer3 = nn.Sequential(
nn.PReLU(),
nn.Conv3d(
in_channels=parameters['output_channels'],
out_channels=parameters['output_channels'],
kernel_size=(parameters['kernel_classification'],
parameters['kernel_classification'],
parameters['kernel_classification']),
stride=parameters['convolution_stride'],
padding=(0, 0, 0)
),
nn.InstanceNorm3d(num_features=parameters['output_channels'])
)
# Other activation functions which might be interesting to test:
# More reading: https://arxiv.org/abs/1706.02515 ; https://mlfromscratch.com/activation-functions-explained/#/
# self.activation = nn.SELU()
# self.activation = nn.ELU()
# self.activation = nn.ReLU()
# Instance normalisation is used to the the small batch size, and as it has shown promise during the experiments with the simple network.
if parameters['dropout'] > 0:
self.dropout_needed = True
self.dropout = nn.Dropout3d(parameters['dropout'])
else:
self.dropout_needed = False
def forward(self, X):
"""Forward pass
Function computing the forward pass through the convolutional layer.
The input to the function is a torch tensor of shape N (batch size) x C (number of channels) x D (input depth) x H (input heigth) x W (input width)
Args:
X (torch.tensor): Input tensor, shape = (N x C x D x H x W)
Returns:
torch.tensor: Output forward passed tensor
"""
feature_map1 = self.convolutional_layer1(X)
feature_map2 = self.convolutional_layer2(feature_map1)