Commit e37a0fe0 authored by Andrei Roibu's avatar Andrei Roibu
Browse files

added dropout for testing in all layers

parent 590e1b32
......@@ -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
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