feat: made model.py extensible

Former-commit-id: 0261a3f6caf571dc19e6ca97d2ecada1e72c7f04

src/unet/model.py

@@ -20,20 +20,24 @@ class UNet(nn.Module):
         self.ups = nn.ModuleList()
         for i in range(len(features) - 1):
             self.ups.append(
-                Up(*features[::-1][i : i + 2]),
+                Up(*features[-1 - i : -1 - i + 3 : -1]),
             )
 
         self.outc = OutConv(features[0], n_classes)
 
     def forward(self, x):
-        x1 = self.inc(x)
-        x2 = self.down1(x1)
-        x3 = self.down2(x2)
-        x4 = self.down3(x3)
-        x5 = self.down4(x4)
-        x = self.up1(x5, x4)
-        x = self.up2(x, x3)
-        x = self.up3(x, x2)
-        x = self.up4(x, x1)
-        logits = self.outc(x)
-        return logits
+
+        skips = []
+
+        x = self.inc(x)
+
+        for down in self.downs:
+            skips.append(x)
+            x = down(x)
+
+        for up, skip in zip(self.ups, reversed(skips)):
+            x = up(x, skip)
+
+        x = self.outc(x)
+
+        return x