update docs for 3D networks

update docs for 3D networks

Author: taigw

pymic/layer/convolution.py

@@ -7,8 +7,22 @@
 class ConvolutionLayer(nn.Module):
     """
     A compose layer with the following components:
-    convolution -> (batch_norm / layer_norm / group_norm / instance_norm) -> activation -> (dropout)
-    batch norm and dropout are optional
+    convolution -> (batch_norm / layer_norm / group_norm / instance_norm) -> (activation) -> (dropout)
+    Batch norm and activation are optional.
+
+    :param in_channels: (int) The input channel number.
+    :param out_channels: (int) The output channel number. 
+    :param kernel_size: The size of convolution kernel. It can be either a single 
+        int or a tupe of two or three ints. 
+    :param dim: (int) The dimention of convolution (2 or 3).
+    :param stride: (int) The stride of convolution. 
+    :param padding: (int) Padding size. 
+    :param dilation: (int) Dilation rate.
+    :param conv_group: (int) The groupt number of convolution. 
+    :param bias: (bool) Add bias or not for convolution. 
+    :param norm_type: (str or None) Normalization type, can be `batch_norm`, 'group_norm'.
+    :param norm_group: (int) The number of group for group normalization.
+    :param acti_func: (str or None) Activation funtion. 
     """
     def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
             stride = 1, padding = 0, dilation = 1, conv_group = 1, bias = True, 
@@ -50,9 +64,23 @@ def forward(self, x):
 
 class DepthSeperableConvolutionLayer(nn.Module):
     """
-    A compose layer with the following components:
-    convolution -> (batch_norm) -> activation -> (dropout)
-    batch norm and dropout are optional
+    Depth seperable convolution with the following components:
+    1x1 conv -> group conv -> (batch_norm / layer_norm / group_norm / instance_norm) -> (activation) -> (dropout)
+    Batch norm and activation are optional.
+
+    :param in_channels: (int) The input channel number.
+    :param out_channels: (int) The output channel number. 
+    :param kernel_size: The size of convolution kernel. It can be either a single 
+        int or a tupe of two or three ints. 
+    :param dim: (int) The dimention of convolution (2 or 3).
+    :param stride: (int) The stride of convolution. 
+    :param padding: (int) Padding size. 
+    :param dilation: (int) Dilation rate.
+    :param conv_group: (int) The groupt number of convolution. 
+    :param bias: (bool) Add bias or not for convolution. 
+    :param norm_type: (str or None) Normalization type, can be `batch_norm`, 'group_norm'.
+    :param norm_group: (int) The number of group for group normalization.
+    :param acti_func: (str or None) Activation funtion. 
     """
     def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
             stride = 1, padding = 0, dilation =1, conv_group = 1, bias = True, 
@@ -97,68 +125,3 @@ def forward(self, x):
             f = self.acti_func(f)
         return f
 
-class ConvolutionSepAll3DLayer(nn.Module):
-    """
-    A compose layer with the following components:
-    convolution -> (batch_norm) -> activation -> (dropout)
-    batch norm and dropout are optional
-    """
-    def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
-            stride = 1, padding = 0, dilation =1, groups = 1, bias = True, 
-            batch_norm = True, acti_func = None):
-        super(ConvolutionSepAll3DLayer, self).__init__()
-        self.n_in_chns  = in_channels
-        self.n_out_chns = out_channels
-        self.batch_norm = batch_norm
-        self.acti_func  = acti_func
-
-        assert(dim == 3)
-        chn = min(in_channels, out_channels)
-
-        self.conv_intra_plane1 = nn.Conv2d(chn, chn,
-            kernel_size, stride, padding, dilation, chn, bias)
-
-        self.conv_intra_plane2 = nn.Conv2d(chn, chn,
-            kernel_size, stride, padding, dilation, chn, bias)
-
-        self.conv_intra_plane3 = nn.Conv2d(chn, chn,
-            kernel_size, stride, padding, dilation, chn, bias)
-
-        self.conv_space_wise  = nn.Conv2d(in_channels, out_channels,
-            1, stride, 0, dilation, 1, bias)
-        
-        if(self.batch_norm):
-                self.bn = nn.BatchNorm3d(out_channels)
-
-    def forward(self, x):
-        in_shape = list(x.shape)
-        assert(len(in_shape) == 5)
-        [B, C, D, H, W] = in_shape
-        f0 = x.permute(0, 2, 1, 3, 4)  #[B, D, C, H, W] 
-        f0 = f0.contiguous().view([B*D, C, H, W])
-
-        Cc = min(self.n_in_chns, self.n_out_chns)
-        Co = self.n_out_chns
-        if(self.n_in_chns > self.n_out_chns):
-            f0 = self.conv_space_wise(f0)  #[B*D, Cc, H, W]  
-   
-        f1 = self.conv_intra_plane1(f0)
-        f2 = f1.contiguous().view([B, D, Cc, H, W])
-        f2 = f2.permute(0, 3, 2, 1, 4) #[B, H, Cc, D, W]
-        f2 = f2.contiguous().view([B*H, Cc, D, W])
-        f2 = self.conv_intra_plane2(f2)
-        f3 = f2.contiguous().view([B, H, Cc, D, W])
-        f3 = f3.permute(0, 4, 2, 3, 1) #[B, W, Cc, D, H]
-        f3 = f3.contiguous().view([B*W, Cc, D, H])
-        f3 = self.conv_intra_plane3(f3)
-        if(self.n_in_chns <= self.n_out_chns):
-            f3 = self.conv_space_wise(f3)  #[B*W, Co, D, H] 
-
-        f3 = f3.contiguous().view([B, W, Co, D, H])
-        f3 = f3.permute([0, 2, 3, 4, 1]) #[B, Co, D, H, W]
-
-        if(self.batch_norm):
-            f3 = self.bn(f3)
-        if(self.acti_func is not None):
-            f3 = self.acti_func(f3)
-        return f3

pymic/layer/deconvolution.py

@@ -7,8 +7,21 @@
 class DeconvolutionLayer(nn.Module):
     """
     A compose layer with the following components:
-    deconvolution -> (batch_norm) -> activation -> (dropout)
-    batch norm and dropout are optional
+    deconvolution -> (batch_norm / layer_norm / group_norm / instance_norm) -> (activation) -> (dropout)
+    Batch norm and activation are optional.
+
+    :param in_channels: (int) The input channel number.
+    :param out_channels: (int) The output channel number. 
+    :param kernel_size: The size of convolution kernel. It can be either a single 
+        int or a tupe of two or three ints. 
+    :param dim: (int) The dimention of convolution (2 or 3).
+    :param stride: (int) The stride of convolution. 
+    :param padding: (int) Padding size. 
+    :param dilation: (int) Dilation rate.
+    :param groups: (int) The groupt number of convolution. 
+    :param bias: (bool) Add bias or not for convolution. 
+    :param batch_norm: (bool) Use batch norm or not.
+    :param acti_func: (str or None) Activation funtion. 
     """
     def __init__(self, in_channels, out_channels, kernel_size, 
             dim = 3, stride = 1, padding = 0, output_padding = 0, 
@@ -44,9 +57,23 @@ def forward(self, x):
 
 class  DepthSeperableDeconvolutionLayer(nn.Module):
     """
-    A compose layer with the following components:
-    convolution -> (batch_norm) -> activation -> (dropout)
-    batch norm and dropout are optional
+    Depth seperable deconvolution with the following components:
+    1x1 conv -> deconv -> (batch_norm / layer_norm / group_norm / instance_norm) -> (activation) -> (dropout)
+    Batch norm and activation are optional.
+
+    :param in_channels: (int) The input channel number.
+    :param out_channels: (int) The output channel number. 
+    :param kernel_size: The size of convolution kernel. It can be either a single 
+        int or a tupe of two or three ints. 
+    :param dim: (int) The dimention of convolution (2 or 3).
+    :param stride: (int) The stride of convolution. 
+    :param padding: (int) Padding size for input.
+    :param output_padding: (int) Padding size for ouput.
+    :param dilation: (int) Dilation rate.
+    :param groups: (int) The groupt number of convolution. 
+    :param bias: (bool) Add bias or not for convolution. 
+    :param batch_norm: (bool) Use batch norm or not.
+    :param acti_func: (str or None) Activation funtion. 
     """
     def __init__(self, in_channels, out_channels, kernel_size, 
         dim = 3, stride = 1, padding = 0, output_padding = 0,

pymic/layer/space2channel.py

@@ -5,7 +5,10 @@
 import torch
 import torch.nn as nn
 import SimpleITK as sitk
+
 class SpaceToChannel3D(nn.Module):
+    """
+    Space to channel transform for 3D input."""
     def __init__(self):
         super(SpaceToChannel3D, self).__init__()
 
@@ -34,6 +37,8 @@ def forward(self, x):
         return x7
 
 class ChannelToSpace3D(nn.Module):
+    """
+    Channel to space transform for 3D input."""
     def __init__(self):
         super(ChannelToSpace3D, self).__init__()
 

pymic/net/cls/torch_pretrained_net.py

@@ -1,6 +1,7 @@
 # pretrained models from pytorch: https://pytorch.org/vision/0.8/models.html
 from __future__ import print_function, division
 
+import itertools
 import torch
 import torch.nn as nn
 import torchvision.models as models
@@ -20,80 +21,149 @@
 #     'mnasnet': models.mnasnet1_0
 # }
 
-class ResNet18(nn.Module):
+class BuiltInNet(nn.Module):
+    """
+    Built-in Network in Pytorch for classification.
+    Parameters should be set in the `params` dictionary that contains the 
+    following fields:
+
+    :param input_chns: (int) Input channel number, default is 3.
+    :param pretrain: (bool) Using pretrained model or not, default is True. 
+    :param update_mode: (str) The strategy for updating layers: "`all`" means updating
+        all the layers, and "`last`" (by default) means updating the last layer, 
+        as well as the first layer when `input_chns` is not 3.
+    """
     def __init__(self, params):
-        super(ResNet18, self).__init__()
-        self.params    = params
-        cls_num  = params['class_num']
-        in_chns  = params.get('input_chns', 3)
+        super(BuiltInNet, self).__init__()
+        self.params   = params
+        self.in_chns  = params.get('input_chns', 3)
         self.pretrain = params.get('pretrain', True)
-        self.update_layers = params.get('update_layers', 0)
+        self.update_mode = params.get('update_mode', "last")
+        self.net = None 
+    
+    def forward(self, x):
+        return self.net(x)
+
+    def get_parameters_to_update(self):
+        pass
+
+class ResNet18(BuiltInNet):
+    """
+    ResNet18 for classification.
+    Parameters should be set in the `params` dictionary that contains the 
+    following fields:
+
+    :param input_chns: (int) Input channel number, default is 3.
+    :param pretrain: (bool) Using pretrained model or not, default is True. 
+    :param update_mode: (str) The strategy for updating layers: "`all`" means updating
+        all the layers, and "`last`" (by default) means updating the last layer, 
+        as well as the first layer when `input_chns` is not 3.
+    """
+    def __init__(self, params):
+        super(ResNet18, self).__init__(params)
         self.net = models.resnet18(pretrained = self.pretrain)
 
         # replace the last layer 
         num_ftrs = self.net.fc.in_features
-        self.net.fc = nn.Linear(num_ftrs, cls_num)
-    
-    def forward(self, x):
-        return self.net(x)
+        self.net.fc = nn.Linear(num_ftrs, params['class_num'])
+
+        # replace the first layer when in_chns is not 3
+        if(self.in_chns != 3):
+            self.net.conv1 = nn.Conv2d(self.in_chns, 64, kernel_size=(7, 7), 
+                stride=(2, 2), padding=(3, 3), bias=False)
 
     def get_parameters_to_update(self):
-        if(self.pretrain == False or self.update_layers == 0):
+        if(self.update_mode == "all"):
             return self.net.parameters()
-        elif(self.update_layers == -1):
-            return self.net.fc.parameters()
+        elif(self.update_layers == "last"):
+            params = self.net.fc.parameters()
+            if(self.in_chns !=3):
+                # combining the two iterables into a single one 
+                # see: https://dzone.com/articles/python-joining-multiple
+                params = itertools.chain()
+                for pram in [self.net.fc.parameters(), self.net.conv1.parameters()]:
+                    params = itertools.chain(params, pram)
+            return  params
         else:
-            raise(ValueError("update_layers can only be 0 (all layers) " +
-                "or -1 (the last layer)"))
+            raise(ValueError("update_mode can only be 'all' or 'last'."))
 
-class VGG16(nn.Module):
+class VGG16(BuiltInNet):
+    """
+    VGG16 for classification.
+    Parameters should be set in the `params` dictionary that contains the 
+    following fields:
+
+    :param input_chns: (int) Input channel number, default is 3.
+    :param pretrain: (bool) Using pretrained model or not, default is True. 
+    :param update_mode: (str) The strategy for updating layers: "`all`" means updating
+        all the layers, and "`last`" (by default) means updating the last layer, 
+        as well as the first layer when `input_chns` is not 3.
+    """
     def __init__(self, params):
-        super(VGG16, self).__init__()
-        self.params    = params
-        cls_num  = params['class_num']
-        in_chns  = params.get('input_chns', 3)
-        self.pretrain = params.get('pretrain', True)
-        self.update_layers = params.get('update_layers', 0)
+        super(VGG16, self).__init__(params)
         self.net = models.vgg16(pretrained = self.pretrain)
 
         # replace the last layer 
         num_ftrs = self.net.classifier[-1].in_features
-        self.net.classifier[-1] = nn.Linear(num_ftrs, cls_num)
-    
-    def forward(self, x):
-        return self.net(x)
+        self.net.classifier[-1] = nn.Linear(num_ftrs, params['class_num'])
+
+        # replace the first layer when in_chns is not 3
+        if(self.in_chns != 3):
+            self.net.features[0] = nn.Conv2d(self.in_chns, 64, kernel_size=(3, 3), 
+                stride=(1, 1), padding=(1, 1), bias=False)
 
     def get_parameters_to_update(self):
-        if(self.pretrain == False or self.update_layers == 0):
+        if(self.update_mode == "all"):
             return self.net.parameters()
-        elif(self.update_layers == -1):
-            return self.net.classifier[-1].parameters()
+        elif(self.update_mode == "last"):
+            params = self.net.classifier[-1].parameters()
+            if(self.in_chns !=3):
+                params = itertools.chain()
+                for pram in [self.net.classifier[-1].parameters(), self.net.net.features[0].parameters()]:
+                    params = itertools.chain(params, pram)
+            return  params
         else:
-            raise(ValueError("update_layers can only be 0 (all layers) " +
-                "or -1 (the last layer)"))
+            raise(ValueError("update_mode can only be 'all' or 'last'."))
+
+class MobileNetV2(BuiltInNet):
+    """
+    MobileNetV2 for classification.
+    Parameters should be set in the `params` dictionary that contains the 
+    following fields:
 
-class MobileNetV2(nn.Module):
+    :param input_chns: (int) Input channel number, default is 3.
+    :param pretrain: (bool) Using pretrained model or not, default is True. 
+    :param update_mode: (str) The strategy for updating layers: "`all`" means updating
+        all the layers, and "`last`" (by default) means updating the last layer, 
+        as well as the first layer when `input_chns` is not 3.
+    """
     def __init__(self, params):
         super(MobileNetV2, self).__init__()
-        self.params = params
-        cls_num  = params['class_num']
-        in_chns  = params.get('input_chns', 3)
-        self.pretrain = params.get('pretrain', True)
-        self.update_layers = params.get('update_layers', 0)
         self.net = models.mobilenet_v2(pretrained = self.pretrain)
 
         # replace the last layer 
         num_ftrs = self.net.last_channel
-        self.net.classifier[-1] = nn.Linear(num_ftrs, cls_num)
-    
-    def forward(self, x):
-        return self.net(x)
+        self.net.classifier[-1] = nn.Linear(num_ftrs, params['class_num'])
+
+        # replace the first layer when in_chns is not 3
+        if(self.in_chns != 3):
+            self.net.features[0][0] = nn.Conv2d(self.in_chns, 32, kernel_size=(3, 3), 
+                stride=(2, 2), padding=(1, 1), bias=False)
 
     def get_parameters_to_update(self):
-        if(self.pretrain == False or self.update_layers == 0):
+        if(self.update_mode == "all"):
             return self.net.parameters()
-        elif(self.update_layers == -1):
-            return self.net.classifier[-1].parameters()
+        elif(self.update_mode == "last"):
+            params = self.net.classifier[-1].parameters()
+            if(self.in_chns !=3):
+                params = itertools.chain()
+                for pram in [self.net.classifier[-1].parameters(), self.net.net.features[0][0].parameters()]:
+                    params = itertools.chain(params, pram)
+            return  params
         else:
-            raise(ValueError("update_layers can only be 0 (all layers) " +
-                "or -1 (the last layer)"))
+            raise(ValueError("update_mode can only be 'all' or 'last'."))
+
+if __name__ == "__main__":
+    params = {"class_num": 2, "pretrain": False, "input_chns": 3}
+    net = ResNet18(params)
+    print(net)