update conv/deconv

Author: taigw

pymic/layer/convolution.py

@@ -28,7 +28,7 @@ def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
                 self.bn = nn.modules.BatchNorm2d(out_channels)
             elif(self.norm_type == 'group_norm'):
                 self.bn = nn.GroupNorm(self.norm_group, out_channels)
-            else:
+            elif(self.norm_type is not None):
                 raise ValueError("unsupported normalization method {0:}".format(norm_type))
         else:        
             self.conv = nn.Conv3d(in_channels, out_channels,
@@ -37,7 +37,7 @@ def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
                 self.bn = nn.modules.BatchNorm3d(out_channels)
             elif(self.norm_type == 'group_norm'):
                 self.bn = nn.GroupNorm(self.norm_group, out_channels)
-            else:
+            elif(self.norm_type is not None):
                 raise ValueError("unsupported normalization method {0:}".format(norm_type))
 
     def forward(self, x):
@@ -74,7 +74,7 @@ def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
                 self.bn = nn.modules.BatchNorm2d(out_channels)
             elif(self.norm_type == 'group_norm'):
                 self.bn = nn.GroupNorm(self.norm_group, out_channels)
-            else:
+            elif(self.norm_type is not None):
                 raise ValueError("unsupported normalization method {0:}".format(norm_type))
         else:        
             self.conv = nn.Conv3d(in_channels, in_channels,
@@ -85,7 +85,7 @@ def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
                 self.bn = nn.modules.BatchNorm3d(out_channels)
             elif(self.norm_type == 'group_norm'):
                 self.bn = nn.GroupNorm(self.norm_group, out_channels)
-            else:
+            elif(self.norm_type is not None):
                 raise ValueError("unsupported normalization method {0:}".format(norm_type))
 
     def forward(self, x):

pymic/layer/deconvolution.py

@@ -41,3 +41,49 @@ def forward(self, x):
         if(self.acti_func is not None):
             f = self.acti_func(f)
         return f
+
+class  DepthSeperableDeconvolutionLayer(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, output_padding = 0,
+        dilation =1, groups = 1, bias = True, 
+        batch_norm = True, acti_func = None):
+        super(DepthSeperableDeconvolutionLayer, self).__init__()
+        self.n_in_chns  = in_channels
+        self.n_out_chns = out_channels
+        self.batch_norm = batch_norm
+        self.acti_func  = acti_func
+        self.groups     = groups
+        assert(dim == 2 or dim == 3)
+        if(dim == 2):
+            self.conv1x1 = nn.Conv2d(in_channels, out_channels,
+                kernel_size = 1, stride = 1, padding = 0, dilation = dilation, 
+                groups = self.groups, bias = bias)
+            self.conv = nn.ConvTranspose2d(out_channels, out_channels,
+                kernel_size, stride, padding, output_padding,
+                groups = out_channels, bias = bias, dilation = dilation)
+            
+            if(self.batch_norm):
+                self.bn = nn.modules.BatchNorm2d(out_channels)
+        else:
+            self.conv1x1 = nn.Conv3d(in_channels, out_channels,
+                kernel_size = 1, stride = 1, padding = 0, dilation = dilation, 
+                groups = self.groups, bias = bias)
+            self.conv = nn.ConvTranspose3d(out_channels, out_channels,
+                kernel_size, stride, padding, output_padding,
+                groups = out_channels, bias = bias, dilation = dilation)
+            if(self.batch_norm):
+                self.bn = nn.modules.BatchNorm3d(out_channels)
+
+    def forward(self, x):
+        f = self.conv1x1(x)
+        f = self.conv(f)
+        if(self.batch_norm):
+            f = self.bn(f)
+        if(self.acti_func is not None):
+            f = self.acti_func(f)
+        return f

pymic/util/rename_model.py

@@ -0,0 +1,21 @@
+import torch
+
+
+def rename_model_variable(input_file, output_file, input_var_list, output_var_list):
+    assert(len(input_var_list) == len(output_var_list))
+    checkpoint = torch.load(input_file)
+    state_dict = checkpoint['model_state_dict']
+    for i in range(len(input_var_list)):
+        input_var  = input_var_list[i]
+        output_var = output_var_list[i]
+        state_dict[output_var] = state_dict[input_var]
+        state_dict.pop(input_var)
+    checkpoint['model_state_dict'] = state_dict
+    torch.save(checkpoint, output_file)
+
+if __name__ == "__main__":
+    input_file  = '/home/disk2t/projects/dlls/training_fetal_brain/model2/unet2dres/model_15000.pt'
+    output_file = '/home/disk2t/projects/dlls/training_fetal_brain/model2/unet2dres/model_15000_rename.pt'
+    input_var_list = ['conv.weight', 'conv.bias']
+    output_var_list= ['conv9.weight', 'conv9.bias']
+    rename_model_variable(input_file, output_file, input_var_list, output_var_list)