From 0c098320c64c07dbb166536ac2681d8194aba9b9 Mon Sep 17 00:00:00 2001 From: Sourcery AI Date: Tue, 5 Oct 2021 10:43:32 +0000 Subject: [PATCH] 'Refactored by Sourcery' --- implementations/aae/aae.py | 14 ++---- implementations/acgan/acgan.py | 5 +- implementations/began/began.py | 7 ++- implementations/bgan/bgan.py | 5 +- implementations/bicyclegan/bicyclegan.py | 5 +- implementations/bicyclegan/models.py | 3 +- implementations/ccgan/ccgan.py | 2 +- implementations/ccgan/models.py | 3 +- implementations/cgan/cgan.py | 5 +- implementations/cluster_gan/clustergan.py | 48 ++++++++----------- implementations/cogan/cogan.py | 2 +- implementations/cogan/mnistm.py | 4 +- .../context_encoder/context_encoder.py | 3 +- implementations/cyclegan/utils.py | 11 ++--- implementations/dcgan/dcgan.py | 9 ++-- implementations/dragan/dragan.py | 12 ++--- implementations/dualgan/dualgan.py | 5 +- implementations/ebgan/ebgan.py | 8 ++-- implementations/esrgan/models.py | 6 ++- implementations/gan/gan.py | 6 +-- implementations/infogan/infogan.py | 5 +- implementations/lsgan/lsgan.py | 9 ++-- implementations/munit/models.py | 22 ++++----- implementations/pix2pix/pix2pix.py | 2 +- implementations/pixelda/mnistm.py | 4 +- implementations/pixelda/pixelda.py | 17 ++----- .../relativistic_gan/relativistic_gan.py | 7 +-- implementations/sgan/sgan.py | 5 +- implementations/softmax_gan/softmax_gan.py | 6 +-- implementations/srgan/models.py | 12 ++--- implementations/stargan/models.py | 8 ++-- implementations/stargan/stargan.py | 3 +- implementations/unit/models.py | 6 +-- implementations/unit/unit.py | 5 +- implementations/wgan/wgan.py | 5 +- implementations/wgan_div/wgan_div.py | 5 +- implementations/wgan_gp/wgan_gp.py | 8 ++-- 37 files changed, 114 insertions(+), 178 deletions(-) diff --git a/implementations/aae/aae.py b/implementations/aae/aae.py index 6bfb0104..9c26f07f 100644 --- a/implementations/aae/aae.py +++ b/implementations/aae/aae.py @@ -33,14 +33,13 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def reparameterization(mu, logvar): std = torch.exp(logvar / 2) sampled_z = Variable(Tensor(np.random.normal(0, 1, (mu.size(0), opt.latent_dim)))) - z = sampled_z * std + mu - return z + return sampled_z * std + mu class Encoder(nn.Module): @@ -63,8 +62,7 @@ def forward(self, img): x = self.model(img_flat) mu = self.mu(x) logvar = self.logvar(x) - z = reparameterization(mu, logvar) - return z + return reparameterization(mu, logvar) class Decoder(nn.Module): @@ -83,8 +81,7 @@ def __init__(self): def forward(self, z): img_flat = self.model(z) - img = img_flat.view(img_flat.shape[0], *img_shape) - return img + return img_flat.view(img_flat.shape[0], *img_shape) class Discriminator(nn.Module): @@ -101,8 +98,7 @@ def __init__(self): ) def forward(self, z): - validity = self.model(z) - return validity + return self.model(z) # Use binary cross-entropy loss diff --git a/implementations/acgan/acgan.py b/implementations/acgan/acgan.py index d7fa2486..dbc28a9d 100644 --- a/implementations/acgan/acgan.py +++ b/implementations/acgan/acgan.py @@ -31,7 +31,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -70,8 +70,7 @@ def forward(self, noise, labels): gen_input = torch.mul(self.label_emb(labels), noise) out = self.l1(gen_input) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): diff --git a/implementations/began/began.py b/implementations/began/began.py index d9d420d8..10332de6 100644 --- a/implementations/began/began.py +++ b/implementations/began/began.py @@ -32,7 +32,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -68,8 +68,7 @@ def __init__(self): def forward(self, noise): out = self.l1(noise) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): @@ -189,7 +188,7 @@ def forward(self, img): diff = torch.mean(gamma * d_loss_real - d_loss_fake) # Update weight term for fake samples - k = k + lambda_k * diff.item() + k += lambda_k * diff.item() k = min(max(k, 0), 1) # Constraint to interval [0, 1] # Update convergence metric diff --git a/implementations/bgan/bgan.py b/implementations/bgan/bgan.py index 4b84eae6..a925c232 100644 --- a/implementations/bgan/bgan.py +++ b/implementations/bgan/bgan.py @@ -34,7 +34,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) class Generator(nn.Module): @@ -78,8 +78,7 @@ def __init__(self): def forward(self, img): img_flat = img.view(img.shape[0], -1) - validity = self.model(img_flat) - return validity + return self.model(img_flat) def boundary_seeking_loss(y_pred, y_true): diff --git a/implementations/bicyclegan/bicyclegan.py b/implementations/bicyclegan/bicyclegan.py index 08c0e04d..e0942a53 100644 --- a/implementations/bicyclegan/bicyclegan.py +++ b/implementations/bicyclegan/bicyclegan.py @@ -45,7 +45,7 @@ os.makedirs("images/%s" % opt.dataset_name, exist_ok=True) os.makedirs("saved_models/%s" % opt.dataset_name, exist_ok=True) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) input_shape = (opt.channels, opt.img_height, opt.img_width) @@ -125,8 +125,7 @@ def sample_images(batches_done): def reparameterization(mu, logvar): std = torch.exp(logvar / 2) sampled_z = Variable(Tensor(np.random.normal(0, 1, (mu.size(0), opt.latent_dim)))) - z = sampled_z * std + mu - return z + return sampled_z * std + mu # ---------- diff --git a/implementations/bicyclegan/models.py b/implementations/bicyclegan/models.py index a5701321..c3ca3c35 100644 --- a/implementations/bicyclegan/models.py +++ b/implementations/bicyclegan/models.py @@ -154,8 +154,7 @@ def discriminator_block(in_filters, out_filters, normalize=True): def compute_loss(self, x, gt): """Computes the MSE between model output and scalar gt""" - loss = sum([torch.mean((out - gt) ** 2) for out in self.forward(x)]) - return loss + return sum(torch.mean((out - gt) ** 2) for out in self.forward(x)) def forward(self, x): outputs = [] diff --git a/implementations/ccgan/ccgan.py b/implementations/ccgan/ccgan.py index 1a55a139..67f05218 100644 --- a/implementations/ccgan/ccgan.py +++ b/implementations/ccgan/ccgan.py @@ -36,7 +36,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) input_shape = (opt.channels, opt.img_size, opt.img_size) diff --git a/implementations/ccgan/models.py b/implementations/ccgan/models.py index 83b9b4fa..bcc2d27b 100644 --- a/implementations/ccgan/models.py +++ b/implementations/ccgan/models.py @@ -38,8 +38,7 @@ 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) - return out + return torch.cat((x, skip_input), 1) class Generator(nn.Module): diff --git a/implementations/cgan/cgan.py b/implementations/cgan/cgan.py index 9e32b7c0..07ab71aa 100644 --- a/implementations/cgan/cgan.py +++ b/implementations/cgan/cgan.py @@ -33,7 +33,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) class Generator(nn.Module): @@ -87,8 +87,7 @@ def __init__(self): def forward(self, img, labels): # Concatenate label embedding and image to produce input d_in = torch.cat((img.view(img.size(0), -1), self.label_embedding(labels)), -1) - validity = self.model(d_in) - return validity + return self.model(d_in) # Loss functions diff --git a/implementations/cluster_gan/clustergan.py b/implementations/cluster_gan/clustergan.py index b1256f12..8bd75fb3 100644 --- a/implementations/cluster_gan/clustergan.py +++ b/implementations/cluster_gan/clustergan.py @@ -105,13 +105,7 @@ def calc_gradient_penalty(netD, real_data, generated_data): # Weight Initializer def initialize_weights(net): for m in net.modules(): - if isinstance(m, nn.Conv2d): - m.weight.data.normal_(0, 0.02) - m.bias.data.zero_() - elif isinstance(m, nn.ConvTranspose2d): - m.weight.data.normal_(0, 0.02) - m.bias.data.zero_() - elif isinstance(m, nn.Linear): + if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear)): m.weight.data.normal_(0, 0.02) m.bias.data.zero_() @@ -292,9 +286,7 @@ def __init__(self, wass_metric=False, verbose=False): print(self.model) def forward(self, img): - # Get output - validity = self.model(img) - return validity + return self.model(img) @@ -323,7 +315,7 @@ def forward(self, img): x_shape = (channels, img_size, img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # Loss function @@ -343,7 +335,7 @@ def forward(self, img): bce_loss.cuda() xe_loss.cuda() mse_loss.cuda() - + Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor # Configure data loader @@ -395,9 +387,11 @@ def forward(self, img): # Training loop print('\nBegin training session with %i epochs...\n'%(n_epochs)) +# Set number of examples for cycle calcs +n_sqrt_samp = 5 for epoch in range(n_epochs): for i, (imgs, itruth_label) in enumerate(dataloader): - + # Ensure generator/encoder are trainable generator.train() encoder.train() @@ -406,16 +400,16 @@ def forward(self, img): generator.zero_grad() encoder.zero_grad() discriminator.zero_grad() - + # Configure input real_imgs = Variable(imgs.type(Tensor)) # --------------------------- # Train Generator + Encoder # --------------------------- - + optimizer_GE.zero_grad() - + # Sample random latent variables zn, zc, zc_idx = sample_z(shape=imgs.shape[0], latent_dim=latent_dim, @@ -423,11 +417,11 @@ def forward(self, img): # Generate a batch of images gen_imgs = generator(zn, zc) - + # Discriminator output from real and generated samples D_gen = discriminator(gen_imgs) D_real = discriminator(real_imgs) - + # Step for Generator & Encoder, n_skip_iter times less than for discriminator if (i % n_skip_iter == 0): # Encode the generated images @@ -463,7 +457,7 @@ def forward(self, img): # Wasserstein GAN loss w/gradient penalty d_loss = torch.mean(D_real) - torch.mean(D_gen) + grad_penalty - + else: # Vanilla GAN loss fake = Variable(Tensor(gen_imgs.size(0), 1).fill_(0.0), requires_grad=False) @@ -484,9 +478,7 @@ def forward(self, img): generator.eval() encoder.eval() - # Set number of examples for cycle calcs - n_sqrt_samp = 5 - n_samp = n_sqrt_samp * n_sqrt_samp + n_samp = n_sqrt_samp**2 ## Cycle through test real -> enc -> gen @@ -499,7 +491,7 @@ def forward(self, img): img_mse_loss = mse_loss(t_imgs, teg_imgs) # Save img reco cycle loss c_i.append(img_mse_loss.item()) - + ## Cycle through randomly sampled encoding -> generator -> encoder zn_samp, zc_samp, zc_samp_idx = sample_z(shape=n_samp, @@ -518,7 +510,7 @@ def forward(self, img): # Save latent space cycle losses c_zn.append(lat_mse_loss.item()) c_zc.append(lat_xe_loss.item()) - + # Save cycled and generated examples! r_imgs, i_label = real_imgs.data[:n_samp], itruth_label[:n_samp] e_zn, e_zc, e_zc_logits = encoder(r_imgs) @@ -529,7 +521,7 @@ def forward(self, img): save_image(gen_imgs_samp.data[:n_samp], 'images/gen_%06i.png' %(epoch), nrow=n_sqrt_samp, normalize=True) - + ## Generate samples for specified classes stack_imgs = [] for idx in range(n_c): @@ -542,7 +534,7 @@ def forward(self, img): # Generate sample instances gen_imgs_samp = generator(zn_samp, zc_samp) - if (len(stack_imgs) == 0): + if not stack_imgs: stack_imgs = gen_imgs_samp else: stack_imgs = torch.cat((stack_imgs, gen_imgs_samp), 0) @@ -551,7 +543,7 @@ def forward(self, img): save_image(stack_imgs, 'images/gen_classes_%06i.png' %(epoch), nrow=n_c, normalize=True) - + print ("[Epoch %d/%d] \n"\ "\tModel Losses: [D: %f] [GE: %f]" % (epoch, @@ -559,7 +551,7 @@ def forward(self, img): d_loss.item(), ge_loss.item()) ) - + print("\tCycle Losses: [x: %f] [z_n: %f] [z_c: %f]"%(img_mse_loss.item(), lat_mse_loss.item(), lat_xe_loss.item()) diff --git a/implementations/cogan/cogan.py b/implementations/cogan/cogan.py index 97b800ce..c213abf5 100644 --- a/implementations/cogan/cogan.py +++ b/implementations/cogan/cogan.py @@ -36,7 +36,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): diff --git a/implementations/cogan/mnistm.py b/implementations/cogan/mnistm.py index e0a9832f..a1de32b3 100644 --- a/implementations/cogan/mnistm.py +++ b/implementations/cogan/mnistm.py @@ -105,9 +105,7 @@ def download(self): os.makedirs(os.path.join(self.root, self.raw_folder)) os.makedirs(os.path.join(self.root, self.processed_folder)) except OSError as e: - if e.errno == errno.EEXIST: - pass - else: + if e.errno != errno.EEXIST: raise # download pkl files diff --git a/implementations/context_encoder/context_encoder.py b/implementations/context_encoder/context_encoder.py index dd828a45..3b66c677 100644 --- a/implementations/context_encoder/context_encoder.py +++ b/implementations/context_encoder/context_encoder.py @@ -8,6 +8,7 @@ 4. Run the sript using command 'python3 context_encoder.py' """ + import argparse import os import numpy as np @@ -46,7 +47,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) # Calculate output of image discriminator (PatchGAN) patch_h, patch_w = int(opt.mask_size / 2 ** 3), int(opt.mask_size / 2 ** 3) diff --git a/implementations/cyclegan/utils.py b/implementations/cyclegan/utils.py index 52c65091..813ff1dc 100644 --- a/implementations/cyclegan/utils.py +++ b/implementations/cyclegan/utils.py @@ -23,13 +23,12 @@ def push_and_pop(self, data): if len(self.data) < self.max_size: self.data.append(element) to_return.append(element) + elif random.uniform(0, 1) > 0.5: + i = random.randint(0, self.max_size - 1) + to_return.append(self.data[i].clone()) + self.data[i] = element else: - if random.uniform(0, 1) > 0.5: - i = random.randint(0, self.max_size - 1) - to_return.append(self.data[i].clone()) - self.data[i] = element - else: - to_return.append(element) + to_return.append(element) return Variable(torch.cat(to_return)) diff --git a/implementations/dcgan/dcgan.py b/implementations/dcgan/dcgan.py index 46e64c06..b7c2a61d 100644 --- a/implementations/dcgan/dcgan.py +++ b/implementations/dcgan/dcgan.py @@ -30,7 +30,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -66,8 +66,7 @@ def __init__(self): def forward(self, z): out = self.l1(z) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): @@ -94,9 +93,7 @@ def discriminator_block(in_filters, out_filters, bn=True): def forward(self, img): out = self.model(img) out = out.view(out.shape[0], -1) - validity = self.adv_layer(out) - - return validity + return self.adv_layer(out) # Loss function diff --git a/implementations/dragan/dragan.py b/implementations/dragan/dragan.py index 6bae1631..c17c7b9c 100644 --- a/implementations/dragan/dragan.py +++ b/implementations/dragan/dragan.py @@ -31,7 +31,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -67,8 +67,7 @@ def __init__(self): def forward(self, noise): out = self.l1(noise) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): @@ -95,9 +94,7 @@ def discriminator_block(in_filters, out_filters, bn=True): def forward(self, img): out = self.model(img) out = out.view(out.shape[0], -1) - validity = self.adv_layer(out) - - return validity + return self.adv_layer(out) # Loss function @@ -163,8 +160,7 @@ def compute_gradient_penalty(D, X): only_inputs=True, )[0] - gradient_penalty = lambda_gp * ((gradients.norm(2, dim=1) - 1) ** 2).mean() - return gradient_penalty + return lambda_gp * ((gradients.norm(2, dim=1) - 1) ** 2).mean() # ---------- diff --git a/implementations/dualgan/dualgan.py b/implementations/dualgan/dualgan.py index e1af8575..12182cec 100644 --- a/implementations/dualgan/dualgan.py +++ b/implementations/dualgan/dualgan.py @@ -47,7 +47,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) # Loss function cycle_loss = torch.nn.L1Loss() @@ -131,8 +131,7 @@ def compute_gradient_penalty(D, real_samples, fake_samples): only_inputs=True, )[0] gradients = gradients.view(gradients.size(0), -1) - gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() - return gradient_penalty + return ((gradients.norm(2, dim=1) - 1) ** 2).mean() def sample_images(batches_done): diff --git a/implementations/ebgan/ebgan.py b/implementations/ebgan/ebgan.py index 0392930b..ef80625d 100644 --- a/implementations/ebgan/ebgan.py +++ b/implementations/ebgan/ebgan.py @@ -32,7 +32,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -67,8 +67,7 @@ def __init__(self): def forward(self, noise): out = self.l1(noise) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): @@ -144,8 +143,7 @@ def pullaway_loss(embeddings): normalized_emb = embeddings / norm similarity = torch.matmul(normalized_emb, normalized_emb.transpose(1, 0)) batch_size = embeddings.size(0) - loss_pt = (torch.sum(similarity) - batch_size) / (batch_size * (batch_size - 1)) - return loss_pt + return (torch.sum(similarity) - batch_size) / (batch_size * (batch_size - 1)) # ---------- diff --git a/implementations/esrgan/models.py b/implementations/esrgan/models.py index ad2c417b..40a0cb8c 100644 --- a/implementations/esrgan/models.py +++ b/implementations/esrgan/models.py @@ -103,8 +103,10 @@ def __init__(self, input_shape): self.output_shape = (1, patch_h, patch_w) def discriminator_block(in_filters, out_filters, first_block=False): - layers = [] - layers.append(nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1)) + layers = [ + nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1) + ] + if not first_block: layers.append(nn.BatchNorm2d(out_filters)) layers.append(nn.LeakyReLU(0.2, inplace=True)) diff --git a/implementations/gan/gan.py b/implementations/gan/gan.py index d6f1d935..7b2338b5 100644 --- a/implementations/gan/gan.py +++ b/implementations/gan/gan.py @@ -32,7 +32,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) class Generator(nn.Module): @@ -76,9 +76,7 @@ def __init__(self): def forward(self, img): img_flat = img.view(img.size(0), -1) - validity = self.model(img_flat) - - return validity + return self.model(img_flat) # Loss function diff --git a/implementations/infogan/infogan.py b/implementations/infogan/infogan.py index ccde076c..e854f146 100644 --- a/implementations/infogan/infogan.py +++ b/implementations/infogan/infogan.py @@ -35,7 +35,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -81,8 +81,7 @@ def forward(self, noise, labels, code): gen_input = torch.cat((noise, labels, code), -1) out = self.l1(gen_input) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): diff --git a/implementations/lsgan/lsgan.py b/implementations/lsgan/lsgan.py index 1930ab84..7872f378 100644 --- a/implementations/lsgan/lsgan.py +++ b/implementations/lsgan/lsgan.py @@ -30,7 +30,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -65,8 +65,7 @@ def __init__(self): def forward(self, z): out = self.l1(z) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): @@ -93,9 +92,7 @@ def discriminator_block(in_filters, out_filters, bn=True): def forward(self, img): out = self.model(img) out = out.view(out.shape[0], -1) - validity = self.adv_layer(out) - - return validity + return self.adv_layer(out) # !!! Minimizes MSE instead of BCE diff --git a/implementations/munit/models.py b/implementations/munit/models.py index e3ce7ec2..81512109 100644 --- a/implementations/munit/models.py +++ b/implementations/munit/models.py @@ -51,12 +51,8 @@ class Decoder(nn.Module): def __init__(self, out_channels=3, dim=64, n_residual=3, n_upsample=2, style_dim=8): super(Decoder, self).__init__() - layers = [] dim = dim * 2 ** n_upsample - # Residual blocks - for _ in range(n_residual): - layers += [ResidualBlock(dim, norm="adain")] - + layers = [ResidualBlock(dim, norm="adain") for _ in range(n_residual)] # Upsampling for _ in range(n_upsample): layers += [ @@ -78,11 +74,11 @@ def __init__(self, out_channels=3, dim=64, n_residual=3, n_upsample=2, style_dim def get_num_adain_params(self): """Return the number of AdaIN parameters needed by the model""" - num_adain_params = 0 - for m in self.modules(): - if m.__class__.__name__ == "AdaptiveInstanceNorm2d": - num_adain_params += 2 * m.num_features - return num_adain_params + return sum( + 2 * m.num_features + for m in self.modules() + if m.__class__.__name__ == "AdaptiveInstanceNorm2d" + ) def assign_adain_params(self, adain_params): """Assign the adain_params to the AdaIN layers in model""" @@ -101,8 +97,7 @@ def assign_adain_params(self, adain_params): def forward(self, content_code, style_code): # Update AdaIN parameters by MLP prediction based off style code self.assign_adain_params(self.mlp(style_code)) - img = self.model(content_code) - return img + return self.model(content_code) ################################# @@ -224,8 +219,7 @@ def discriminator_block(in_filters, out_filters, normalize=True): def compute_loss(self, x, gt): """Computes the MSE between model output and scalar gt""" - loss = sum([torch.mean((out - gt) ** 2) for out in self.forward(x)]) - return loss + return sum(torch.mean((out - gt) ** 2) for out in self.forward(x)) def forward(self, x): outputs = [] diff --git a/implementations/pix2pix/pix2pix.py b/implementations/pix2pix/pix2pix.py index dda43e77..e470106f 100644 --- a/implementations/pix2pix/pix2pix.py +++ b/implementations/pix2pix/pix2pix.py @@ -44,7 +44,7 @@ os.makedirs("images/%s" % opt.dataset_name, exist_ok=True) os.makedirs("saved_models/%s" % opt.dataset_name, exist_ok=True) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) # Loss functions criterion_GAN = torch.nn.MSELoss() diff --git a/implementations/pixelda/mnistm.py b/implementations/pixelda/mnistm.py index b55b29d9..f923a345 100644 --- a/implementations/pixelda/mnistm.py +++ b/implementations/pixelda/mnistm.py @@ -115,9 +115,7 @@ def download(self): os.makedirs(os.path.join(self.root, self.raw_folder)) os.makedirs(os.path.join(self.root, self.processed_folder)) except OSError as e: - if e.errno == errno.EEXIST: - pass - else: + if e.errno != errno.EEXIST: raise # download pkl files diff --git a/implementations/pixelda/pixelda.py b/implementations/pixelda/pixelda.py index eb80b2d0..902b83a7 100644 --- a/implementations/pixelda/pixelda.py +++ b/implementations/pixelda/pixelda.py @@ -39,7 +39,7 @@ patch = int(opt.img_size / 2 ** 4) patch = (1, patch, patch) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -76,9 +76,7 @@ def __init__(self): self.l1 = nn.Sequential(nn.Conv2d(opt.channels * 2, 64, 3, 1, 1), nn.ReLU(inplace=True)) - resblocks = [] - for _ in range(opt.n_residual_blocks): - resblocks.append(ResidualBlock()) + resblocks = [ResidualBlock() for _ in range(opt.n_residual_blocks)] self.resblocks = nn.Sequential(*resblocks) self.l2 = nn.Sequential(nn.Conv2d(64, opt.channels, 3, 1, 1), nn.Tanh()) @@ -87,9 +85,7 @@ def forward(self, img, z): gen_input = torch.cat((img, self.fc(z).view(*img.shape)), 1) out = self.l1(gen_input) out = self.resblocks(out) - img_ = self.l2(out) - - return img_ + return self.l2(out) class Discriminator(nn.Module): @@ -112,9 +108,7 @@ def block(in_features, out_features, normalization=True): ) def forward(self, img): - validity = self.model(img) - - return validity + return self.model(img) class Classifier(nn.Module): @@ -138,8 +132,7 @@ def block(in_features, out_features, normalization=True): def forward(self, img): feature_repr = self.model(img) feature_repr = feature_repr.view(feature_repr.size(0), -1) - label = self.output_layer(feature_repr) - return label + return self.output_layer(feature_repr) # Loss function diff --git a/implementations/relativistic_gan/relativistic_gan.py b/implementations/relativistic_gan/relativistic_gan.py index 0447bdb5..fa234444 100644 --- a/implementations/relativistic_gan/relativistic_gan.py +++ b/implementations/relativistic_gan/relativistic_gan.py @@ -58,8 +58,7 @@ def __init__(self): def forward(self, z): out = self.l1(z) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): @@ -86,9 +85,7 @@ def discriminator_block(in_filters, out_filters, bn=True): def forward(self, img): out = self.model(img) out = out.view(out.shape[0], -1) - validity = self.adv_layer(out) - - return validity + return self.adv_layer(out) # Loss function diff --git a/implementations/sgan/sgan.py b/implementations/sgan/sgan.py index d12bf292..8321e2a1 100644 --- a/implementations/sgan/sgan.py +++ b/implementations/sgan/sgan.py @@ -31,7 +31,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) def weights_init_normal(m): @@ -69,8 +69,7 @@ def __init__(self): def forward(self, noise): out = self.l1(noise) out = out.view(out.shape[0], 128, self.init_size, self.init_size) - img = self.conv_blocks(out) - return img + return self.conv_blocks(out) class Discriminator(nn.Module): diff --git a/implementations/softmax_gan/softmax_gan.py b/implementations/softmax_gan/softmax_gan.py index 23a0775b..fad5e2e5 100644 --- a/implementations/softmax_gan/softmax_gan.py +++ b/implementations/softmax_gan/softmax_gan.py @@ -32,7 +32,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) class Generator(nn.Module): @@ -75,9 +75,7 @@ def __init__(self): def forward(self, img): img_flat = img.view(img.shape[0], -1) - validity = self.model(img_flat) - - return validity + return self.model(img_flat) # Loss function diff --git a/implementations/srgan/models.py b/implementations/srgan/models.py index b8efd3c0..d79e113c 100644 --- a/implementations/srgan/models.py +++ b/implementations/srgan/models.py @@ -38,9 +38,7 @@ def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=16): self.conv1 = nn.Sequential(nn.Conv2d(in_channels, 64, kernel_size=9, stride=1, padding=4), nn.PReLU()) # Residual blocks - res_blocks = [] - for _ in range(n_residual_blocks): - res_blocks.append(ResidualBlock(64)) + res_blocks = [ResidualBlock(64) for _ in range(n_residual_blocks)] self.res_blocks = nn.Sequential(*res_blocks) # Second conv layer post residual blocks @@ -48,7 +46,7 @@ def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=16): # Upsampling layers upsampling = [] - for out_features in range(2): + for _ in range(2): upsampling += [ # nn.Upsample(scale_factor=2), nn.Conv2d(64, 256, 3, 1, 1), @@ -81,8 +79,10 @@ def __init__(self, input_shape): self.output_shape = (1, patch_h, patch_w) def discriminator_block(in_filters, out_filters, first_block=False): - layers = [] - layers.append(nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1)) + layers = [ + nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1) + ] + if not first_block: layers.append(nn.BatchNorm2d(out_filters)) layers.append(nn.LeakyReLU(0.2, inplace=True)) diff --git a/implementations/stargan/models.py b/implementations/stargan/models.py index ce339e3a..42b3bd6f 100644 --- a/implementations/stargan/models.py +++ b/implementations/stargan/models.py @@ -65,7 +65,7 @@ def __init__(self, img_shape=(3, 128, 128), res_blocks=9, c_dim=5): nn.InstanceNorm2d(curr_dim // 2, affine=True, track_running_stats=True), nn.ReLU(inplace=True), ] - curr_dim = curr_dim // 2 + curr_dim //= 2 # Output layer model += [nn.Conv2d(curr_dim, channels, 7, stride=1, padding=3), nn.Tanh()] @@ -91,8 +91,10 @@ def __init__(self, img_shape=(3, 128, 128), c_dim=5, n_strided=6): def discriminator_block(in_filters, out_filters): """Returns downsampling layers of each discriminator block""" - layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1), nn.LeakyReLU(0.01)] - return layers + return [ + nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1), + nn.LeakyReLU(0.01), + ] layers = discriminator_block(channels, 64) curr_dim = 64 diff --git a/implementations/stargan/stargan.py b/implementations/stargan/stargan.py index e2244800..de7de015 100644 --- a/implementations/stargan/stargan.py +++ b/implementations/stargan/stargan.py @@ -157,8 +157,7 @@ def compute_gradient_penalty(D, real_samples, fake_samples): only_inputs=True, )[0] gradients = gradients.view(gradients.size(0), -1) - gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() - return gradient_penalty + return ((gradients.norm(2, dim=1) - 1) ** 2).mean() label_changes = [ diff --git a/implementations/unit/models.py b/implementations/unit/models.py index 10c5e9d3..9d002f83 100644 --- a/implementations/unit/models.py +++ b/implementations/unit/models.py @@ -96,12 +96,8 @@ def __init__(self, out_channels=3, dim=64, n_upsample=2, shared_block=None): self.shared_block = shared_block - layers = [] dim = dim * 2 ** n_upsample - # Residual blocks - for _ in range(3): - layers += [ResidualBlock(dim)] - + layers = [ResidualBlock(dim) for _ in range(3)] # Upsampling for _ in range(n_upsample): layers += [ diff --git a/implementations/unit/unit.py b/implementations/unit/unit.py index 30b88774..c224c03c 100644 --- a/implementations/unit/unit.py +++ b/implementations/unit/unit.py @@ -41,7 +41,7 @@ opt = parser.parse_args() print(opt) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) # Create sample and checkpoint directories os.makedirs("images/%s" % opt.dataset_name, exist_ok=True) @@ -162,8 +162,7 @@ def sample_images(batches_done): def compute_kl(mu): mu_2 = torch.pow(mu, 2) - loss = torch.mean(mu_2) - return loss + return torch.mean(mu_2) # ---------- diff --git a/implementations/wgan/wgan.py b/implementations/wgan/wgan.py index 6c9008a7..3e042439 100644 --- a/implementations/wgan/wgan.py +++ b/implementations/wgan/wgan.py @@ -33,7 +33,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) class Generator(nn.Module): @@ -76,8 +76,7 @@ def __init__(self): def forward(self, img): img_flat = img.view(img.shape[0], -1) - validity = self.model(img_flat) - return validity + return self.model(img_flat) # Initialize generator and discriminator diff --git a/implementations/wgan_div/wgan_div.py b/implementations/wgan_div/wgan_div.py index 1b0a8096..e2ae17f3 100644 --- a/implementations/wgan_div/wgan_div.py +++ b/implementations/wgan_div/wgan_div.py @@ -36,7 +36,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) class Generator(nn.Module): @@ -79,8 +79,7 @@ def __init__(self): def forward(self, img): img_flat = img.view(img.shape[0], -1) - validity = self.model(img_flat) - return validity + return self.model(img_flat) k = 2 diff --git a/implementations/wgan_gp/wgan_gp.py b/implementations/wgan_gp/wgan_gp.py index 17d5afbf..2a8211a8 100644 --- a/implementations/wgan_gp/wgan_gp.py +++ b/implementations/wgan_gp/wgan_gp.py @@ -36,7 +36,7 @@ img_shape = (opt.channels, opt.img_size, opt.img_size) -cuda = True if torch.cuda.is_available() else False +cuda = bool(torch.cuda.is_available()) class Generator(nn.Module): @@ -79,8 +79,7 @@ def __init__(self): def forward(self, img): img_flat = img.view(img.shape[0], -1) - validity = self.model(img_flat) - return validity + return self.model(img_flat) # Loss weight for gradient penalty @@ -134,8 +133,7 @@ def compute_gradient_penalty(D, real_samples, fake_samples): only_inputs=True, )[0] gradients = gradients.view(gradients.size(0), -1) - gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() - return gradient_penalty + return ((gradients.norm(2, dim=1) - 1) ** 2).mean() # ----------