Keep patchGAN output shape in discriminators

Author: eriklindernoren

Diff for: implementations/bicyclegan/bicyclegan.py

@@ -46,13 +46,13 @@
 
 cuda = True if torch.cuda.is_available() else False
 
-img_shape = (opt.channels, opt.img_height, opt.img_width)
+input_shape = (opt.channels, opt.img_height, opt.img_width)
 
 # Loss functions
 mae_loss = torch.nn.L1Loss()
 
 # Initialize generator, encoder and discriminators
-generator = Generator(opt.latent_dim, img_shape)
+generator = Generator(opt.latent_dim, input_shape)
 encoder = Encoder(opt.latent_dim)
 D_VAE = MultiDiscriminator()
 D_LR = MultiDiscriminator()

Diff for: implementations/ccgan/ccgan.py

@@ -18,44 +18,34 @@
 import torch.nn.functional as F
 import torch
 
-os.makedirs('images', exist_ok=True)
+os.makedirs("images", exist_ok=True)
 
 parser = argparse.ArgumentParser()
-parser.add_argument('--n_epochs', type=int, default=200, help='number of epochs of training')
-parser.add_argument('--batch_size', type=int, default=8, help='size of the batches')
-parser.add_argument('--dataset_name', type=str, default='img_align_celeba', help='name of the dataset')
-parser.add_argument('--lr', type=float, default=0.0002, help='adam: learning rate')
-parser.add_argument('--b1', type=float, default=0.5, help='adam: decay of first order momentum of gradient')
-parser.add_argument('--b2', type=float, default=0.999, help='adam: decay of first order momentum of gradient')
-parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads to use during batch generation')
-parser.add_argument('--latent_dim', type=int, default=100, help='dimensionality of the latent space')
-parser.add_argument('--img_size', type=int, default=128, help='size of each image dimension')
-parser.add_argument('--mask_size', type=int, default=32, help='size of random mask')
-parser.add_argument('--channels', type=int, default=3, help='number of image channels')
-parser.add_argument('--sample_interval', type=int, default=500, help='interval between image sampling')
+parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
+parser.add_argument("--batch_size", type=int, default=8, help="size of the batches")
+parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
+parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
+parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
+parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
+parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
+parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
+parser.add_argument("--img_size", type=int, default=128, help="size of each image dimension")
+parser.add_argument("--mask_size", type=int, default=32, help="size of random mask")
+parser.add_argument("--channels", type=int, default=3, help="number of image channels")
+parser.add_argument("--sample_interval", type=int, default=500, help="interval between image sampling")
 opt = parser.parse_args()
 print(opt)
 
 cuda = True if torch.cuda.is_available() else False
 
-# Calculate output of image discriminator (PatchGAN)
-patch_h, patch_w = int(opt.img_size / 2**3), int(opt.img_size / 2**3)
-patch = (1, patch_h, patch_w)
-
-def weights_init_normal(m):
-    classname = m.__class__.__name__
-    if classname.find('Conv') != -1:
-        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
-    elif classname.find('BatchNorm2d') != -1:
-        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
-        torch.nn.init.constant_(m.bias.data, 0.0)
+input_shape = (opt.channels, opt.img_size, opt.img_size)
 
 # Loss function
 adversarial_loss = torch.nn.MSELoss()
 
 # Initialize generator and discriminator
-generator = Generator(channels=opt.channels)
-discriminator = Discriminator(channels=opt.channels)
+generator = Generator(input_shape)
+discriminator = Discriminator(input_shape)
 
 if cuda:
     generator.cuda()
@@ -67,28 +57,32 @@ def weights_init_normal(m):
 discriminator.apply(weights_init_normal)
 
 # Dataset loader
-transforms_ = [ transforms.Resize((opt.img_size, opt.img_size), Image.BICUBIC),
-                transforms.ToTensor(),
-                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)) ]
-transforms_lr = [ transforms.Resize((opt.img_size//4, opt.img_size//4), Image.BICUBIC),
-                transforms.ToTensor(),
-                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)) ]
-dataloader = DataLoader(ImageDataset("../../data/%s" % opt.dataset_name,
-                        transforms_x=transforms_, transforms_lr=transforms_lr),
-                        batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_cpu)
+transforms_ = [
+    transforms.Resize((opt.img_size, opt.img_size), Image.BICUBIC),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+]
+transforms_lr = [
+    transforms.Resize((opt.img_size // 4, opt.img_size // 4), Image.BICUBIC),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+]
+dataloader = DataLoader(
+    ImageDataset("../../data/%s" % opt.dataset_name, transforms_x=transforms_, transforms_lr=transforms_lr),
+    batch_size=opt.batch_size,
+    shuffle=True,
+    num_workers=opt.n_cpu,
+)
 
 # Optimizers
 optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
 optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
 
 Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
 
-# Adversarial ground truths
-valid = Variable(Tensor(np.ones(patch)), requires_grad=False)
-fake = Variable(Tensor(np.zeros(patch)), requires_grad=False)
 
 def apply_random_mask(imgs):
-    idx = np.random.randint(0, opt.img_size-opt.mask_size, (imgs.shape[0], 2))
+    idx = np.random.randint(0, opt.img_size - opt.mask_size, (imgs.shape[0], 2))
 
     masked_imgs = imgs.clone()
     for i, (y1, x1) in enumerate(idx):
@@ -97,25 +91,26 @@ def apply_random_mask(imgs):
 
     return masked_imgs
 
+
 def save_sample(saved_samples):
     # Generate inpainted image
-    gen_imgs = generator(saved_samples['masked'], saved_samples['lowres'])
+    gen_imgs = generator(saved_samples["masked"], saved_samples["lowres"])
     # Save sample
-    sample = torch.cat((saved_samples['masked'].data, gen_imgs.data, saved_samples['imgs'].data), -2)
-    save_image(sample,'images/%d.png' % batches_done, nrow=5, normalize=True)
+    sample = torch.cat((saved_samples["masked"].data, gen_imgs.data, saved_samples["imgs"].data), -2)
+    save_image(sample, "images/%d.png" % batches_done, nrow=5, normalize=True)
 
 
 saved_samples = {}
 for epoch in range(opt.n_epochs):
     for i, batch in enumerate(dataloader):
-        imgs = batch['x']
-        imgs_lr = batch['x_lr']
+        imgs = batch["x"]
+        imgs_lr = batch["x_lr"]
 
         masked_imgs = apply_random_mask(imgs)
 
         # Adversarial ground truths
-        valid = Variable(Tensor(imgs.shape[0], *patch).fill_(1.0), requires_grad=False)
-        fake = Variable(Tensor(imgs.shape[0], *patch).fill_(0.0), requires_grad=False)
+        valid = Variable(Tensor(imgs.shape[0], *discriminator.output_shape).fill_(1.0), requires_grad=False)
+        fake = Variable(Tensor(imgs.shape[0], *discriminator.output_shape).fill_(0.0), requires_grad=False)
 
         if cuda:
             imgs = imgs.type(Tensor)
@@ -155,18 +150,20 @@ def save_sample(saved_samples):
         d_loss.backward()
         optimizer_D.step()
 
-        print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" % (epoch, opt.n_epochs, i, len(dataloader),
-                                                            d_loss.item(), g_loss.item()))
+        print(
+            "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
+            % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())
+        )
 
         # Save first ten samples
         if not saved_samples:
-            saved_samples['imgs'] = real_imgs[:1].clone()
-            saved_samples['masked'] = masked_imgs[:1].clone()
-            saved_samples['lowres'] = imgs_lr[:1].clone()
-        elif saved_samples['imgs'].size(0) < 10:
-            saved_samples['imgs'] = torch.cat((saved_samples['imgs'], real_imgs[:1]), 0)
-            saved_samples['masked'] = torch.cat((saved_samples['masked'], masked_imgs[:1]), 0)
-            saved_samples['lowres'] = torch.cat((saved_samples['lowres'], imgs_lr[:1]), 0)
+            saved_samples["imgs"] = real_imgs[:1].clone()
+            saved_samples["masked"] = masked_imgs[:1].clone()
+            saved_samples["lowres"] = imgs_lr[:1].clone()
+        elif saved_samples["imgs"].size(0) < 10:
+            saved_samples["imgs"] = torch.cat((saved_samples["imgs"], real_imgs[:1]), 0)
+            saved_samples["masked"] = torch.cat((saved_samples["masked"], masked_imgs[:1]), 0)
+            saved_samples["lowres"] = torch.cat((saved_samples["lowres"], imgs_lr[:1]), 0)
 
         batches_done = epoch * len(dataloader) + i
         if batches_done % opt.sample_interval == 0:

Diff for: implementations/ccgan/models.py

@@ -6,6 +6,7 @@
 #           U-NET
 ##############################
 
+
 class UNetDown(nn.Module):
     def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
         super(UNetDown, self).__init__()
@@ -21,12 +22,15 @@ def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
     def forward(self, x):
         return self.model(x)
 
+
 class UNetUp(nn.Module):
     def __init__(self, in_size, out_size, dropout=0.0):
         super(UNetUp, self).__init__()
-        model = [   nn.ConvTranspose2d(in_size, out_size, 4, stride=2, padding=1, bias=False),
-                    nn.BatchNorm2d(out_size, 0.8),
-                    nn.ReLU(inplace=True)]
+        model = [
+            nn.ConvTranspose2d(in_size, out_size, 4, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(out_size, 0.8),
+            nn.ReLU(inplace=True),
+        ]
         if dropout:
             model.append(nn.Dropout(dropout))
 
@@ -35,16 +39,16 @@ def __init__(self, in_size, out_size, dropout=0.0):
     def forward(self, x, skip_input):
         x = self.model(x)
         out = torch.cat((x, skip_input), 1)
-        #out = torch.add(x, skip_input)
         return out
 
+
 class Generator(nn.Module):
-    def __init__(self, channels=3):
+    def __init__(self, input_shape):
         super(Generator, self).__init__()
-
+        channels, _, _ = input_shape
         self.down1 = UNetDown(channels, 64, normalize=False)
         self.down2 = UNetDown(64, 128)
-        self.down3 = UNetDown(128+channels, 256, dropout=0.5)
+        self.down3 = UNetDown(128 + channels, 256, dropout=0.5)
         self.down4 = UNetDown(256, 512, dropout=0.5)
         self.down5 = UNetDown(512, 512, dropout=0.5)
         self.down6 = UNetDown(512, 512, dropout=0.5)
@@ -53,12 +57,9 @@ def __init__(self, channels=3):
         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+channels, 64)
+        self.up5 = UNetUp(256 + channels, 64)
 
-
-        final = [   nn.Upsample(scale_factor=2),
-                    nn.Conv2d(128, channels, 3, 1, 1),
-                    nn.Tanh() ]
+        final = [nn.Upsample(scale_factor=2), nn.Conv2d(128, channels, 3, 1, 1), nn.Tanh()]
         self.final = nn.Sequential(*final)
 
     def forward(self, x, x_lr):
@@ -78,10 +79,16 @@ def forward(self, x, x_lr):
 
         return self.final(u5)
 
+
 class Discriminator(nn.Module):
-    def __init__(self, channels=3):
+    def __init__(self, input_shape):
         super(Discriminator, self).__init__()
 
+        channels, height, width = input_shape
+        # Calculate output of image discriminator (PatchGAN)
+        patch_h, patch_w = int(height / 2 ** 3), int(width / 2 ** 3)
+        self.output_shape = (1, patch_h, patch_w)
+
         def discriminator_block(in_filters, out_filters, stride, normalize):
             """Returns layers of each discriminator block"""
             layers = [nn.Conv2d(in_filters, out_filters, 3, stride, 1)]
@@ -92,10 +99,7 @@ def discriminator_block(in_filters, out_filters, stride, normalize):
 
         layers = []
         in_filters = channels
-        for out_filters, stride, normalize in [ (64, 2, False),
-                                                (128, 2, True),
-                                                (256, 2, True),
-                                                (512, 1, True)]:
+        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
 

Diff for: implementations/cgan/cgan.py

@@ -93,7 +93,6 @@ def forward(self, img, labels):
 
 # Loss functions
 adversarial_loss = torch.nn.MSELoss()
-auxiliary_loss = torch.nn.CrossEntropyLoss()
 
 # Initialize generator and discriminator
 generator = Generator()
@@ -103,7 +102,6 @@ def forward(self, img, labels):
     generator.cuda()
     discriminator.cuda()
     adversarial_loss.cuda()
-    auxiliary_loss.cuda()
 
 # Configure data loader
 os.makedirs("../../data/mnist", exist_ok=True)