modules.py 12.1 KB
Newer Older
1
2
3
4
"""Deep Learning Modules

Description:

5
    This folder contains several the building blocks for the translation neural network.
6

7
8
9
10
11
12
Usage:

    To use the modules, import the packages and instantiate any module/block class as you wish:

        from utils.modules import modules as module_names
        block = module_name.ConvolutionalBlock(parameters)
13
14
15
16
17
18
19

"""

import torch
import torch.nn as nn
import torch.nn.functional as F

Andrei Roibu's avatar
Andrei Roibu committed
20
21
# 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!

22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

# 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
        """

95
96
97
98
99
100
        X = self.convolutional_layer(X)

        if self.dropout_needed:
            X = self.dropout(X)

        return X
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186


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
        """

187
188
189
190
191
192
193
        X = torch.add(self.convolutional_layer2(
            self.convolutional_layer(X)), X)

        if self.dropout_needed:
            X = self.dropout(X)

        return X
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228


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(
229
230
231
232
233
234
235
236
            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'])
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
        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
        """

259
260
        X = self.activation(self.normalization(
            self.transpose_convolutional_layer(X, output_size=Y_encoder_size)))
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308

        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)
        )
309
310
        self.normalization = nn.InstanceNorm3d(
            num_features=parameters['number_of_classes'])
311

312
313
314
315
316
317
        if parameters['final_activation'] == 'sigmoid':
            self.activation = nn.Sigmoid()
        elif parameters['final_activation'] == 'tanh':
            self.activation = nn.Tanh()
        else:
            self.activation = None
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335

        # 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))

336
337
        if isinstance(self.activation, (nn.Sigmoid, nn.Tanh)):
            logits = self.activation(logits)
338

339
        return logits