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

added modules for AE and scaling for UNet

parent 9ca11f63
......@@ -19,6 +19,306 @@ 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!
# CycleGAN 3D Generator Autoencoder:
class ResNetEncoderBlock3D(nn.Module):
"""Parent class for a 3D convolutional block.
This class represents a generic parent class for a convolutional 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(ResNetEncoderBlock3D, 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_heigth'] - 1) / 2)
padding_depth = int((parameters['kernel_heigth'] - 1) / 2)
self.convolutional_layer = nn.Sequential(
nn.Conv3d(
in_channels=parameters['input_channels'],
out_channels=parameters['output_channels'],
kernel_size=parameters['kernel_heigth'],
stride=parameters['convolution_stride'],
padding=(padding_depth, padding_heigth, padding_width)
),
nn.InstanceNorm3d(num_features=parameters['output_channels']),
nn.PReLU(),
)
# 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
"""
return self.convolutional_layer(X)
class ResNetBlock3D(nn.Module):
"""Parent class for a 3D ResNet convolutional block.
This class represents a generic parent class for a residual convolutional 3D 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(ResNetBlock3D, 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)
self.convolutional_layer = nn.Sequential(
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']),
nn.PReLU(),
)
self.convolutional_layer2 = nn.Sequential(
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'])
)
# 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
"""
return torch.add(self.convolutional_layer2(self.convolutional_layer(X)), X)
class ResNetDecoderBlock3D(nn.Module):
"""Forward 3D decoder path block for a CycleGAN Generator.
This class creates a simple decoder block using transpose convolutions
Args:
parameters (dict): Contains information relevant parameters
parameters = {
'kernel_heigth': 5
'kernel_width': 5
'kernel_depth': 5
'input_channels': 64
'output_channels': 64
'convolution_stride': 1
'dropout': 0.2
'pool_kernel_size': 2
'pool_stride': 2
'up_mode': 'upconv'
}
Returns:
Y (torch.tensor): Output forward passed tensor through the decoder block
"""
def __init__(self, parameters):
super(ResNetDecoderBlock3D, self).__init__()
padding_heigth = int((parameters['pool_kernel_size'] - 1) / 2)
padding_width = int((parameters['pool_kernel_size'] - 1) / 2)
padding_depth = int((parameters['pool_kernel_size'] - 1) / 2)
self.transpose_convolutional_layer = nn.ConvTranspose3d(
in_channels=parameters['input_channels'],
out_channels=parameters['output_channels'],
kernel_size=parameters['pool_kernel_size'],
stride=parameters['pool_stride'],
padding=(padding_depth, padding_heigth, padding_width)
)
self.normalization = nn.InstanceNorm3d(num_features=parameters['output_channels'])
self.activation = nn.PReLU()
if parameters['dropout'] > 0:
self.dropout_needed = True
self.dropout = nn.Dropout3d(parameters['dropout'])
else:
self.dropout_needed = False
def forward(self, X, Y_encoder_size):
"""Forward pass for ResNet decoder block
Function computing the forward pass through the decoder block.
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)
Y_encoder_size (torch.tensor): Shape of the corresponding tensor from the encoder path, required to ensure that the dimensions are kept consistent
Returns:
X (torch.tensor): Output forward passed tensor through the decoder block
"""
X = self.activation(self.normalization(self.transpose_convolutional_layer(X, output_size=Y_encoder_size)))
if self.dropout_needed:
X = self.dropout(X)
return X
class ResNetClassifierBlock3D(nn.Module):
"""Classifier block for a CGan Autoencoder Generator.
This class creates a simple classifier block following the architecture:
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': 1
'convolution_stride': 1
'dropout': 0.2
'pool_kernel_size': 2
'pool_stride': 2
'up_mode': 'upconv'
'number_of_classes': 1
}
Returns:
Y (torch.tensor): Output forward passed tensor through the decoder block
"""
def __init__(self, parameters):
super(ResNetClassifierBlock3D, self).__init__()
padding_heigth = int((parameters['kernel_classification'] - 1) / 2)
padding_width = int((parameters['kernel_classification'] - 1) / 2)
padding_depth = int((parameters['kernel_classification'] - 1) / 2)
self.convolutional_layer = nn.Conv3d(
in_channels=parameters['input_channels'],
out_channels=parameters['number_of_classes'],
kernel_size=parameters['kernel_classification'],
stride=parameters['convolution_stride'],
padding=(padding_depth, padding_heigth, padding_width)
)
self.normalization = nn.InstanceNorm3d(num_features=parameters['number_of_classes'])
# self.activation = nn.Sigmoid()
self.activation = nn.Tanh()
# TODO: Might be wworth looking at GANS for image generation, and adding padding
def forward(self, X):
"""Forward pass for U-net classifier block
Function computing the forward pass through the classifier block.
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:
logits (torch.tensor): Output logits from forward pass tensor through the classifier block
"""
logits = self.normalization(self.convolutional_layer(X))
logits = self.activation(logits)
return logits
# Classical 3D UNet:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment