feat(semseg): allow model customization (#1371)

* feat(semantic_segmentation): allow customization of unet

* feat(semseg): allow model customization

* style(semseg): format to PEP8

* fix(semseg): rename logger

* docs(changelog): updated semantic segmentation example

* suggestions

* suggestions

* flake8

Co-authored-by: J. Borovec <jirka.borovec@seznam.cz>

pl_examples/models/unet.py

@@ -9,39 +9,46 @@ class UNet(nn.Module):
     Link - https://arxiv.org/abs/1505.04597
 
     Parameters:
-        num_classes (int): Number of output classes required (default 19 for KITTI dataset)
-        bilinear (bool): Whether to use bilinear interpolation or transposed
+        num_classes: Number of output classes required (default 19 for KITTI dataset)
+        num_layers: Number of layers in each side of U-net
+        features_start: Number of features in first layer
+        bilinear: Whether to use bilinear interpolation or transposed
             convolutions for upsampling.
     """
 
-    def __init__(self, num_classes=19, bilinear=False):
+    def __init__(
+            self, num_classes: int = 19,
+            num_layers: int = 5,
+            features_start: int = 64,
+            bilinear: bool = False
+    ):
         super().__init__()
-        self.layer1 = DoubleConv(3, 64)
-        self.layer2 = Down(64, 128)
-        self.layer3 = Down(128, 256)
-        self.layer4 = Down(256, 512)
-        self.layer5 = Down(512, 1024)
+        self.num_layers = num_layers
 
-        self.layer6 = Up(1024, 512, bilinear=bilinear)
-        self.layer7 = Up(512, 256, bilinear=bilinear)
-        self.layer8 = Up(256, 128, bilinear=bilinear)
-        self.layer9 = Up(128, 64, bilinear=bilinear)
+        layers = [DoubleConv(3, features_start)]
 
-        self.layer10 = nn.Conv2d(64, num_classes, kernel_size=1)
+        feats = features_start
+        for _ in range(num_layers - 1):
+            layers.append(Down(feats, feats * 2))
+            feats *= 2
+
+        for _ in range(num_layers - 1):
+            layers.append(Up(feats, feats // 2), bilinear)
+            feats //= 2
+
+        layers.append(nn.Conv2d(feats, num_classes, kernel_size=1))
+
+        self.layers = nn.ModuleList(layers)
 
     def forward(self, x):
-        x1 = self.layer1(x)
-        x2 = self.layer2(x1)
-        x3 = self.layer3(x2)
-        x4 = self.layer4(x3)
-        x5 = self.layer5(x4)
-
-        x6 = self.layer6(x5, x4)
-        x6 = self.layer7(x6, x3)
-        x6 = self.layer8(x6, x2)
-        x6 = self.layer9(x6, x1)
-
-        return self.layer10(x6)
+        xi = [self.layers[0](x)]
+        # Down path
+        for layer in self.layers[1:self.num_layers]:
+            xi.append(layer(xi[-1]))
+        # Up path
+        for i, layer in enumerate(self.layers[self.num_layers:-1]):
+            xi[-1] = layer(xi[-1], xi[-2 - i])
+        return self.layers[-1](xi[-1])
 
 
 class DoubleConv(nn.Module):
@@ -50,7 +57,7 @@ class DoubleConv(nn.Module):
     (3x3 conv -> BN -> ReLU) ** 2
     """
 
-    def __init__(self, in_ch, out_ch):
+    def __init__(self, in_ch: int, out_ch: int):
         super().__init__()
         self.net = nn.Sequential(
             nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
@@ -70,7 +77,7 @@ class Down(nn.Module):
     Combination of MaxPool2d and DoubleConv in series
     """
 
-    def __init__(self, in_ch, out_ch):
+    def __init__(self, in_ch: int, out_ch: int):
         super().__init__()
         self.net = nn.Sequential(
             nn.MaxPool2d(kernel_size=2, stride=2),
@@ -88,7 +95,7 @@ class Up(nn.Module):
     followed by double 3x3 convolution.
     """
 
-    def __init__(self, in_ch, out_ch, bilinear=False):
+    def __init__(self, in_ch: int, out_ch: int, bilinear: bool = False):
         super().__init__()
         self.upsample = None
         if bilinear: