added dropout for testing in all layers

utils/modules.py

@@ -92,7 +92,12 @@ class ResNetEncoderBlock3D(nn.Module):
             torch.tensor: Output forward passed tensor
         """
 
-        return self.convolutional_layer(X)
+        X = self.convolutional_layer(X)
+
+        if self.dropout_needed:
+            X = self.dropout(X)
+
+        return X
 
 
 class ResNetBlock3D(nn.Module):
@@ -179,7 +184,13 @@ class ResNetBlock3D(nn.Module):
             torch.tensor: Output forward passed tensor
         """
 
-        return torch.add(self.convolutional_layer2(self.convolutional_layer(X)), X)
+        X = torch.add(self.convolutional_layer2(
+            self.convolutional_layer(X)), X)
+
+        if self.dropout_needed:
+            X = self.dropout(X)
+
+        return X
 
 
 class ResNetDecoderBlock3D(nn.Module):
@@ -215,13 +226,14 @@ class ResNetDecoderBlock3D(nn.Module):
         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'])
+            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:
@@ -244,7 +256,8 @@ class ResNetDecoderBlock3D(nn.Module):
             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)))
+        X = self.activation(self.normalization(
+            self.transpose_convolutional_layer(X, output_size=Y_encoder_size)))
 
         if self.dropout_needed:
             X = self.dropout(X)
@@ -293,7 +306,8 @@ class ResNetClassifierBlock3D(nn.Module):
             stride=parameters['convolution_stride'],
             padding=(padding_depth, padding_heigth, padding_width)
         )
-        self.normalization = nn.InstanceNorm3d(num_features=parameters['number_of_classes'])
+        self.normalization = nn.InstanceNorm3d(
+            num_features=parameters['number_of_classes'])
 
         # self.activation = nn.Sigmoid()
         self.activation = nn.Tanh()
@@ -317,4 +331,4 @@ class ResNetClassifierBlock3D(nn.Module):
 
         logits = self.activation(logits)
 
-        return logits
\ No newline at end of file
+        return logits