readds Identity instead of using lambda

Author: dnlcrl

src/models/classifiers_2D.py

@@ -1,5 +1,5 @@
 import torch.nn as nn
-from models.model_utils import conv2d_bn_block, dense_layer_bn
+from models.model_utils import conv2d_bn_block, dense_layer_bn, Identity
 
 
 class NormalNet2D(nn.Module):
@@ -27,7 +27,7 @@ def __init__(self, n_channels=1, nlabels=1, init_filters=32):
         )
         self.classifier = nn.Sequential(
             dense_layer_bn(16*nf, 16*nf),
-            dense_layer_bn(16*nf, nlabels, activation=(lambda x: x))
+            dense_layer_bn(16*nf, nlabels, activation=Identity)
         )
 
     def forward(self, x):
@@ -61,7 +61,7 @@ def __init__(self, n_channels=1, nlabels=1, init_filters=32):
         )
         self.classifier = nn.Sequential(
             nn.AvgPool2d(2),
-            dense_layer_bn(16*nf, nlabels, activation=lambda x: x)
+            dense_layer_bn(16*nf, nlabels, activation=Identity)
         )
 
     def forward(self, x):
@@ -98,7 +98,7 @@ def __init__(self, n_channels=1, nlabels=1, init_filters=32):
         )
         self.classifier = nn.Sequential(
             dense_layer_bn(16*nf, 16*nf),
-            dense_layer_bn(16*nf, nlabels, activation=lambda x: x)
+            dense_layer_bn(16*nf, nlabels, activation=Identity)
         )
 
     def forward(self, x):
@@ -130,7 +130,7 @@ def __init__(self, n_channels=1, nlabels=1, init_filters=32):
         )
         self.classifier = nn.Sequential(
             nn.AvgPool2d(2),
-            dense_layer_bn(16*nf, nlabels, activation=lambda x: x)
+            dense_layer_bn(16*nf, nlabels, activation=Identity)
         )
 
     def forward(self, x):

src/models/classifiers_3D.py

@@ -1,5 +1,5 @@
 import torch.nn as nn
-from models.model_utils import conv3d_bn_block, dense_layer_bn
+from models.model_utils import conv3d_bn_block, dense_layer_bn, Identity
 
 
 class FCNBN(nn.Module):
@@ -26,7 +26,7 @@ def __init__(self, n_channels=1, nlabels=1, init_filters=32):
             conv3d_bn_block(8*nf, 8*nf),
         )
         self.classifier = nn.Sequential(
-            conv3d_bn_block(8*nf, nlabels, kernel=1, activation=lambda x: x),
+            conv3d_bn_block(8*nf, nlabels, kernel=1, activation=Identity),
             nn.AvgPool3d(2),
         )
 
@@ -56,7 +56,7 @@ def __init__(self, n_channels=1, nlabels=1, init_filters=32):
             conv3d_bn_block(8*nf, 8*nf),
         )
         self.classifier = nn.Sequential(
-            conv3d_bn_block(8*nf, nlabels, kernel=1, activation=lambda x: x),
+            conv3d_bn_block(8*nf, nlabels, kernel=1, activation=Identity),
             nn.AvgPool3d(2),
         )
 
@@ -88,7 +88,7 @@ def __init__(self, n_channels=1, nlabels=1, init_filters=32):
         )
         self.classifier = nn.Sequential(
             dense_layer_bn(8*nf, 16*nf),
-            dense_layer_bn(16*nf, nlabels, activation=lambda x: x),
+            dense_layer_bn(16*nf, nlabels, activation=Identity),
         )
 
     def forward(self, x):

src/models/mask_generators.py

@@ -75,4 +75,3 @@ def forward(self, x):
         cat1 = crop_and_concat(xu1, x0)
         x7 = self.conv7(cat1)
         return x7
-

src/models/model_utils.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-ACTIVATION = nn.ReLU
+ACTIVATION = nn.LeakyReLU
 
 
 def crop_and_concat(upsampled, bypass, crop=False):