added modules for AE and scaling for UNet

utils/modules.py

@@ -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: