Code cleanup

Author: eriklindernoren

figures/discogan.png

@@ -146,6 +146,10 @@ def __init__(self, in_channels=3, out_channels=3, res_blocks=9):
     def forward(self, x):
         return self.model(x)
 
+##############################
+#        Discriminator
+##############################
+
 class Discriminator(nn.Module):
     def __init__(self, in_channels=3):
         super(Discriminator, self).__init__()
@@ -158,18 +162,13 @@ def discriminator_block(in_filters, out_filters, stride, normalize):
             layers.append(nn.LeakyReLU(0.2, inplace=True))
             return layers
 
-        layers = []
-        in_filters = in_channels
-        for out_filters, stride, normalize in [ (64, 2, False),
-                                                (128, 2, True),
-                                                (256, 2, True),
-                                                (512, 1, True)]:
-            layers.extend(discriminator_block(in_filters, out_filters, stride, normalize))
-            in_filters = out_filters
-
-        layers.append(nn.Conv2d(out_filters, 1, 3, 1, 1))
-
-        self.model = nn.Sequential(*layers)
+        self.model = nn.Sequential(
+            *discriminator_block(in_channels, 64, 2, False),
+            *discriminator_block(64, 128, 2, True),
+            *discriminator_block(128, 256, 2, True),
+            *discriminator_block(256, 512, 1, True),
+            nn.Conv2d(512, 1, 3, 1, 1)
+        )
 
     def forward(self, img):
         return self.model(img)

implementations/cyclegan/models.py

@@ -185,4 +185,4 @@ def forward(self, img):
 
         batches_done = epoch * len(dataloader) + i
         if batches_done % opt.sample_interval == 0:
-            save_image(gen_imgs.data[:25], 'images/%d.png' % batches_done, nrow=5, normalize=True)
+            save_image(gen_imgs.data[:100], 'images/%d.png' % batches_done, nrow=10, normalize=True)

implementations/dcgan/dcgan.py

@@ -54,12 +54,14 @@ def __init__(self, in_channels=3, out_channels=3):
         self.down4 = UNetDown(256, 512, dropout=0.5)
         self.down5 = UNetDown(512, 512, dropout=0.5)
         self.down6 = UNetDown(512, 512, dropout=0.5)
+        self.down7 = UNetDown(512, 512, dropout=0.5)
 
-        self.up1 = UNetUp(512, 512, dropout=0.5)
+        self.up1 = UNetUp(512, 512, dropout=0.8)
         self.up2 = UNetUp(1024, 512, dropout=0.5)
-        self.up3 = UNetUp(1024, 256, dropout=0.5)
-        self.up4 = UNetUp(512, 128)
-        self.up5 = UNetUp(256, 64)
+        self.up3 = UNetUp(1024, 512, dropout=0.5)
+        self.up4 = UNetUp(1024, 256, dropout=0.5)
+        self.up5 = UNetUp(512, 128)
+        self.up6 = UNetUp(256, 64)
 
 
         final = [   nn.Upsample(scale_factor=2),
@@ -75,14 +77,19 @@ def forward(self, x):
         d4 = self.down4(d3)
         d5 = self.down5(d4)
         d6 = self.down6(d5)
-        u1 = self.up1(d6, d5)
-        u2 = self.up2(u1, d4)
-        u3 = self.up3(u2, d3)
-        u4 = self.up4(u3, d2)
-        u5 = self.up5(u4, d1)
+        d7 = self.down7(d6)
+        u1 = self.up1(d7, d6)
+        u2 = self.up2(u1, d5)
+        u3 = self.up3(u2, d4)
+        u4 = self.up4(u3, d3)
+        u5 = self.up5(u4, d2)
+        u6 = self.up6(u5, d1)
 
-        return self.final(u5)
+        return self.final(u6)
 
+##############################
+#        Discriminator
+##############################
 
 class Discriminator(nn.Module):
     def __init__(self, img_shape):
@@ -104,12 +111,5 @@ def block(in_features, out_features, normalization=True):
             nn.Conv2d(512, 1, 3, 1, 1)
         )
 
-        #in_size = img_shape[1] // 2**4
-        #self.output_layer = nn.Sequential(nn.Linear(512*in_size**2, 1))
-
     def forward(self, img):
-        # feature_repr = self.model(img)
-        # feature_repr = feature_repr.view(feature_repr.size(0), -1)
-        # validity = self.output_layer(feature_repr)
-
         return self.model(img)

implementations/dualgan/models.py

@@ -98,7 +98,7 @@ def forward(self, img):
 
         return validity
 
-# Use binary cross-entropy loss
+# !!! Minimizes MSE instead of BCE
 adversarial_loss = torch.nn.MSELoss()
 
 # Initialize generator and discriminator

implementations/lsgan/lsgan.py

@@ -150,6 +150,10 @@ def __init__(self, in_channels=3, out_channels=3, resblocks=9):
     def forward(self, x):
         return self.model(x)
 
+##############################
+#        Discriminator
+##############################
+
 class Discriminator(nn.Module):
     def __init__(self, in_channels=3):
         super(Discriminator, self).__init__()
@@ -162,19 +166,13 @@ def discriminator_block(in_filters, out_filters, stride, normalize):
             layers.append(nn.LeakyReLU(0.2, inplace=True))
             return layers
 
-        layers = []
-        in_filters = in_channels*2
-        for out_filters, stride, normalize in [ (64, 2, False),
-                                                (128, 2, True),
-                                                (256, 2, True),
-                                                (512, 2, True)]:
-            layers.extend(discriminator_block(in_filters, out_filters, stride, normalize))
-            in_filters = out_filters
-
-        # Output layer
-        layers.append(nn.Conv2d(out_filters, 1, 3, 1, 1))
-
-        self.model = nn.Sequential(*layers)
+        self.model = nn.Sequential(
+            *discriminator_block(in_channels, 64, 2, False),
+            *discriminator_block(64, 128, 2, True),
+            *discriminator_block(128, 256, 2, True),
+            *discriminator_block(256, 512, 2, True),
+            nn.Conv2d(512, 1, 3, 1, 1)
+        )
 
     def forward(self, img_A, img_B):
         # Concatenate image and condition image by channels to produce input