From c5d6be170013db8975d96806084df573d87e5d26 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Linder-Nor=C3=A9n?= Date: Thu, 28 Mar 2019 15:55:25 +0100 Subject: [PATCH] MNIST normalization. Black refactoring. --- .gitignore | 1 + implementations/aae/aae.py | 75 ++++++----- implementations/acgan/acgan.py | 92 +++++++------ implementations/began/began.py | 91 +++++++------ implementations/bgan/bgan.py | 61 +++++---- implementations/cgan/cgan.py | 64 +++++---- implementations/cogan/cogan.py | 109 ++++++++------- implementations/cogan/mnistm.py | 83 +++++------- implementations/dcgan/dcgan.py | 71 +++++----- implementations/dragan/dragan.py | 88 +++++++----- implementations/ebgan/ebgan.py | 89 ++++++------ implementations/gan/gan.py | 54 +++++--- implementations/infogan/infogan.py | 128 ++++++++++-------- implementations/lsgan/lsgan.py | 70 +++++----- implementations/pixelda/pixelda.py | 150 ++++++++++++--------- implementations/sgan/sgan.py | 85 ++++++------ implementations/softmax_gan/softmax_gan.py | 64 +++++---- implementations/wgan/wgan.py | 55 ++++---- implementations/wgan_div/wgan_div.py | 34 +++-- implementations/wgan_gp/wgan_gp.py | 4 +- 20 files changed, 801 insertions(+), 667 deletions(-) diff --git a/.gitignore b/.gitignore index 58b8b40d..9fc40d1d 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ .DS_Store data/*/ +implementations/*/data implementations/*/images implementations/*/saved_models diff --git a/implementations/aae/aae.py b/implementations/aae/aae.py index 405e9df6..6bfb0104 100644 --- a/implementations/aae/aae.py +++ b/implementations/aae/aae.py @@ -15,19 +15,19 @@ 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=64, help='size of the batches') -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=10, help='dimensionality of the latent code') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, 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=64, help="size of the batches") +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=10, help="dimensionality of the latent code") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling") opt = parser.parse_args() print(opt) @@ -35,12 +35,14 @@ cuda = True if torch.cuda.is_available() else False + 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 + class Encoder(nn.Module): def __init__(self): super(Encoder, self).__init__() @@ -50,7 +52,7 @@ def __init__(self): nn.LeakyReLU(0.2, inplace=True), nn.Linear(512, 512), nn.BatchNorm1d(512), - nn.LeakyReLU(0.2, inplace=True) + nn.LeakyReLU(0.2, inplace=True), ) self.mu = nn.Linear(512, opt.latent_dim) @@ -64,6 +66,7 @@ def forward(self, img): z = reparameterization(mu, logvar) return z + class Decoder(nn.Module): def __init__(self): super(Decoder, self).__init__() @@ -75,7 +78,7 @@ def __init__(self): nn.BatchNorm1d(512), nn.LeakyReLU(0.2, inplace=True), nn.Linear(512, int(np.prod(img_shape))), - nn.Tanh() + nn.Tanh(), ) def forward(self, z): @@ -83,6 +86,7 @@ def forward(self, z): img = img_flat.view(img_flat.shape[0], *img_shape) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() @@ -93,13 +97,14 @@ def __init__(self): nn.Linear(512, 256), nn.LeakyReLU(0.2, inplace=True), nn.Linear(256, 1), - nn.Sigmoid() + nn.Sigmoid(), ) def forward(self, z): validity = self.model(z) return validity + # Use binary cross-entropy loss adversarial_loss = torch.nn.BCELoss() pixelwise_loss = torch.nn.L1Loss() @@ -117,29 +122,36 @@ def forward(self, z): pixelwise_loss.cuda() # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers -optimizer_G = torch.optim.Adam( itertools.chain(encoder.parameters(), decoder.parameters()), - lr=opt.lr, betas=(opt.b1, opt.b2)) +optimizer_G = torch.optim.Adam( + itertools.chain(encoder.parameters(), decoder.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 + def sample_image(n_row, batches_done): """Saves a grid of generated digits""" # Sample noise - z = Variable(Tensor(np.random.normal(0, 1, (n_row**2, opt.latent_dim)))) + z = Variable(Tensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim)))) gen_imgs = decoder(z) - save_image(gen_imgs.data, 'images/%d.png' % batches_done, nrow=n_row, normalize=True) + save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=n_row, normalize=True) + # ---------- # Training @@ -165,8 +177,9 @@ def sample_image(n_row, batches_done): decoded_imgs = decoder(encoded_imgs) # Loss measures generator's ability to fool the discriminator - g_loss = 0.001 * adversarial_loss(discriminator(encoded_imgs), valid) + \ - 0.999 * pixelwise_loss(decoded_imgs, real_imgs) + g_loss = 0.001 * adversarial_loss(discriminator(encoded_imgs), valid) + 0.999 * pixelwise_loss( + decoded_imgs, real_imgs + ) g_loss.backward() optimizer_G.step() @@ -188,8 +201,10 @@ def sample_image(n_row, batches_done): 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()) + ) batches_done = epoch * len(dataloader) + i if batches_done % opt.sample_interval == 0: diff --git a/implementations/acgan/acgan.py b/implementations/acgan/acgan.py index b0c99e71..d7fa2486 100644 --- a/implementations/acgan/acgan.py +++ b/implementations/acgan/acgan.py @@ -14,41 +14,43 @@ 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=64, help='size of the batches') -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('--n_classes', type=int, default=10, help='number of classes for dataset') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, 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=64, help="size of the batches") +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("--n_classes", type=int, default=10, help="number of classes for dataset") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling") opt = parser.parse_args() print(opt) cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm2d') != -1: + 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) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.label_emb = nn.Embedding(opt.n_classes, opt.latent_dim) - self.init_size = opt.img_size // 4 # Initial size before upsampling - self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.init_size = opt.img_size // 4 # Initial size before upsampling + self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.BatchNorm2d(128), @@ -61,7 +63,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, noise, labels): @@ -71,15 +73,14 @@ def forward(self, noise, labels): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() def discriminator_block(in_filters, out_filters, bn=True): """Returns layers of each discriminator block""" - block = [ nn.Conv2d(in_filters, out_filters, 3, 2, 1), - nn.LeakyReLU(0.2, inplace=True), - nn.Dropout2d(0.25)] + block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)] if bn: block.append(nn.BatchNorm2d(out_filters, 0.8)) return block @@ -92,13 +93,11 @@ def discriminator_block(in_filters, out_filters, bn=True): ) # The height and width of downsampled image - ds_size = opt.img_size // 2**4 + ds_size = opt.img_size // 2 ** 4 # Output layers - self.adv_layer = nn.Sequential( nn.Linear(128*ds_size**2, 1), - nn.Sigmoid()) - self.aux_layer = nn.Sequential( nn.Linear(128*ds_size**2, opt.n_classes), - nn.Softmax()) + self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid()) + self.aux_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.n_classes), nn.Softmax()) def forward(self, img): out = self.conv_blocks(img) @@ -108,6 +107,7 @@ def forward(self, img): return validity, label + # Loss functions adversarial_loss = torch.nn.BCELoss() auxiliary_loss = torch.nn.CrossEntropyLoss() @@ -127,15 +127,19 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -144,15 +148,17 @@ def forward(self, img): FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor + def sample_image(n_row, batches_done): """Saves a grid of generated digits ranging from 0 to n_classes""" # Sample noise - z = Variable(FloatTensor(np.random.normal(0, 1, (n_row**2, opt.latent_dim)))) + z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim)))) # Get labels ranging from 0 to n_classes for n rows labels = np.array([num for _ in range(n_row) for num in range(n_row)]) labels = Variable(LongTensor(labels)) gen_imgs = generator(z, labels) - save_image(gen_imgs.data, 'images/%d.png' % batches_done, nrow=n_row, normalize=True) + save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=n_row, normalize=True) + # ---------- # Training @@ -186,8 +192,7 @@ def sample_image(n_row, batches_done): # Loss measures generator's ability to fool the discriminator validity, pred_label = discriminator(gen_imgs) - g_loss = 0.5 * (adversarial_loss(validity, valid) + \ - auxiliary_loss(pred_label, gen_labels)) + g_loss = 0.5 * (adversarial_loss(validity, valid) + auxiliary_loss(pred_label, gen_labels)) g_loss.backward() optimizer_G.step() @@ -200,13 +205,11 @@ def sample_image(n_row, batches_done): # Loss for real images real_pred, real_aux = discriminator(real_imgs) - d_real_loss = (adversarial_loss(real_pred, valid) + \ - auxiliary_loss(real_aux, labels)) / 2 + d_real_loss = (adversarial_loss(real_pred, valid) + auxiliary_loss(real_aux, labels)) / 2 # Loss for fake images fake_pred, fake_aux = discriminator(gen_imgs.detach()) - d_fake_loss = (adversarial_loss(fake_pred, fake) + \ - auxiliary_loss(fake_aux, gen_labels)) / 2 + d_fake_loss = (adversarial_loss(fake_pred, fake) + auxiliary_loss(fake_aux, gen_labels)) / 2 # Total discriminator loss d_loss = (d_real_loss + d_fake_loss) / 2 @@ -219,9 +222,10 @@ def sample_image(n_row, batches_done): d_loss.backward() optimizer_D.step() - print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]" % (epoch, opt.n_epochs, i, len(dataloader), - d_loss.item(), 100 * d_acc, - g_loss.item())) + print( + "[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]" + % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), 100 * d_acc, g_loss.item()) + ) batches_done = epoch * len(dataloader) + i if batches_done % opt.sample_interval == 0: sample_image(n_row=10, batches_done=batches_done) diff --git a/implementations/began/began.py b/implementations/began/began.py index 6b329e5c..d9d420d8 100644 --- a/implementations/began/began.py +++ b/implementations/began/began.py @@ -14,19 +14,19 @@ 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=64, help='size of the batches') -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=62, help='dimensionality of the latent space') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, help='number of image channels') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +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=62, help="dimensionality of the latent space") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="number of image channels") opt = parser.parse_args() print(opt) @@ -34,20 +34,22 @@ cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm2d') != -1: + 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) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.init_size = opt.img_size // 4 - self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.BatchNorm2d(128), @@ -60,7 +62,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, noise): @@ -69,31 +71,26 @@ def forward(self, noise): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() # Upsampling - self.down = nn.Sequential( - nn.Conv2d(opt.channels, 64, 3, 2, 1), - nn.ReLU(), - ) + self.down = nn.Sequential(nn.Conv2d(opt.channels, 64, 3, 2, 1), nn.ReLU()) # Fully-connected layers - self.down_size = (opt.img_size // 2) - down_dim = 64 * (opt.img_size // 2)**2 + self.down_size = opt.img_size // 2 + down_dim = 64 * (opt.img_size // 2) ** 2 self.fc = nn.Sequential( nn.Linear(down_dim, 32), nn.BatchNorm1d(32, 0.8), nn.ReLU(inplace=True), nn.Linear(32, down_dim), nn.BatchNorm1d(down_dim), - nn.ReLU(inplace=True) + nn.ReLU(inplace=True), ) # Upsampling - self.up = nn.Sequential( - nn.Upsample(scale_factor=2), - nn.Conv2d(64, opt.channels, 3, 1, 1) - ) + self.up = nn.Sequential(nn.Upsample(scale_factor=2), nn.Conv2d(64, opt.channels, 3, 1, 1)) def forward(self, img): out = self.down(img) @@ -101,6 +98,7 @@ def forward(self, img): out = self.up(out.view(out.size(0), 64, self.down_size, self.down_size)) return out + # Initialize generator and discriminator generator = Generator() discriminator = Discriminator() @@ -114,15 +112,19 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -137,7 +139,7 @@ def forward(self, img): # BEGAN hyper parameters gamma = 0.75 lambda_k = 0.001 -k = 0. +k = 0.0 for epoch in range(opt.n_epochs): for i, (imgs, _) in enumerate(dataloader): @@ -180,27 +182,28 @@ def forward(self, img): d_loss.backward() optimizer_D.step() - #---------------- + # ---------------- # Update weights - #---------------- + # ---------------- diff = torch.mean(gamma * d_loss_real - d_loss_fake) # Update weight term for fake samples k = k + lambda_k * diff.item() - k = min(max(k, 0), 1) # Constraint to interval [0, 1] + k = min(max(k, 0), 1) # Constraint to interval [0, 1] # Update convergence metric M = (d_loss_real + torch.abs(diff)).data[0] - #-------------- + # -------------- # Log Progress - #-------------- + # -------------- - print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] -- M: %f, k: %f" % (epoch, opt.n_epochs, i, len(dataloader), - d_loss.item(), g_loss.item(), - M, k)) + print( + "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] -- M: %f, k: %f" + % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item(), M, k) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/bgan/bgan.py b/implementations/bgan/bgan.py index 8d9895ca..4b84eae6 100644 --- a/implementations/bgan/bgan.py +++ b/implementations/bgan/bgan.py @@ -16,19 +16,19 @@ 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=64, help='size of the batches') -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=28, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, help='interval betwen image samples') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +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=28, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples") opt = parser.parse_args() print(opt) @@ -36,12 +36,13 @@ cuda = True if torch.cuda.is_available() else False + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() def block(in_feat, out_feat, normalize=True): - layers = [ nn.Linear(in_feat, out_feat)] + layers = [nn.Linear(in_feat, out_feat)] if normalize: layers.append(nn.BatchNorm1d(out_feat, 0.8)) layers.append(nn.LeakyReLU(0.2, inplace=True)) @@ -61,6 +62,7 @@ def forward(self, z): img = img.view(img.shape[0], *img_shape) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() @@ -71,7 +73,7 @@ def __init__(self): nn.Linear(512, 256), nn.LeakyReLU(0.2, inplace=True), nn.Linear(256, 1), - nn.Sigmoid() + nn.Sigmoid(), ) def forward(self, img): @@ -79,12 +81,14 @@ def forward(self, img): validity = self.model(img_flat) return validity + def boundary_seeking_loss(y_pred, y_true): """ Boundary seeking loss. Reference: https://wiseodd.github.io/techblog/2017/03/07/boundary-seeking-gan/ """ - return 0.5 * torch.mean((torch.log(y_pred) - torch.log(1 - y_pred))**2) + return 0.5 * torch.mean((torch.log(y_pred) - torch.log(1 - y_pred)) ** 2) + discriminator_loss = torch.nn.BCELoss() @@ -98,14 +102,19 @@ def boundary_seeking_loss(y_pred, y_true): discriminator_loss.cuda() # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) -mnist_loader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) +os.makedirs("../../data/mnist", exist_ok=True) +dataloader = torch.utils.data.DataLoader( + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -155,9 +164,11 @@ def boundary_seeking_loss(y_pred, y_true): 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(mnist_loader), - 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(mnist_loader), d_loss.item(), g_loss.item()) + ) batches_done = epoch * len(mnist_loader) + 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/cgan/cgan.py b/implementations/cgan/cgan.py index 645cbb98..780f82bf 100644 --- a/implementations/cgan/cgan.py +++ b/implementations/cgan/cgan.py @@ -14,20 +14,20 @@ 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=64, help='size of the batches') -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('--n_classes', type=int, default=10, help='number of classes for dataset') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, 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=64, help="size of the batches") +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("--n_classes", type=int, default=10, help="number of classes for dataset") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling") opt = parser.parse_args() print(opt) @@ -35,6 +35,7 @@ cuda = True if torch.cuda.is_available() else False + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() @@ -42,14 +43,14 @@ def __init__(self): self.label_emb = nn.Embedding(opt.n_classes, opt.n_classes) def block(in_feat, out_feat, normalize=True): - layers = [ nn.Linear(in_feat, out_feat)] + layers = [nn.Linear(in_feat, out_feat)] if normalize: layers.append(nn.BatchNorm1d(out_feat, 0.8)) layers.append(nn.LeakyReLU(0.2, inplace=True)) return layers self.model = nn.Sequential( - *block(opt.latent_dim+opt.n_classes, 128, normalize=False), + *block(opt.latent_dim + opt.n_classes, 128, normalize=False), *block(128, 256), *block(256, 512), *block(512, 1024), @@ -64,6 +65,7 @@ def forward(self, noise, labels): img = img.view(img.size(0), *img_shape) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() @@ -79,7 +81,7 @@ def __init__(self): nn.Linear(512, 512), nn.Dropout(0.4), nn.LeakyReLU(0.2, inplace=True), - nn.Linear(512, 1) + nn.Linear(512, 1), ) def forward(self, img, labels): @@ -88,6 +90,7 @@ def forward(self, img, labels): validity = self.model(d_in) return validity + # Loss functions adversarial_loss = torch.nn.MSELoss() auxiliary_loss = torch.nn.CrossEntropyLoss() @@ -103,15 +106,19 @@ def forward(self, img, labels): auxiliary_loss.cuda() # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -124,12 +131,13 @@ def forward(self, img, labels): def sample_image(n_row, batches_done): """Saves a grid of generated digits ranging from 0 to n_classes""" # Sample noise - z = Variable(FloatTensor(np.random.normal(0, 1, (n_row**2, opt.latent_dim)))) + z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim)))) # Get labels ranging from 0 to n_classes for n rows labels = np.array([num for _ in range(n_row) for num in range(n_row)]) labels = Variable(LongTensor(labels)) gen_imgs = generator(z, labels) - save_image(gen_imgs.data, 'images/%d.png' % batches_done, nrow=n_row, normalize=True) + save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=n_row, normalize=True) + # ---------- # Training @@ -188,8 +196,10 @@ def sample_image(n_row, batches_done): 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()) + ) batches_done = epoch * len(dataloader) + i if batches_done % opt.sample_interval == 0: diff --git a/implementations/cogan/cogan.py b/implementations/cogan/cogan.py index 545ea2ee..97b800ce 100644 --- a/implementations/cogan/cogan.py +++ b/implementations/cogan/cogan.py @@ -18,19 +18,19 @@ 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=32, help='size of the batches') -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=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=3, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, help='interval betwen image samples') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=32, help="size of the batches") +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=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=3, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples") opt = parser.parse_args() print(opt) @@ -38,20 +38,22 @@ cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Linear') != -1: + if classname.find("Linear") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm') != -1: + elif classname.find("BatchNorm") != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0) + class CoupledGenerators(nn.Module): def __init__(self): super(CoupledGenerators, self).__init__() self.init_size = opt.img_size // 4 - self.fc = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.fc = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.shared_conv = nn.Sequential( nn.BatchNorm2d(128), @@ -66,14 +68,14 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) self.G2 = nn.Sequential( nn.Conv2d(128, 64, 3, stride=1, padding=1), nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, noise): @@ -84,6 +86,7 @@ def forward(self, noise): img2 = self.G2(img_emb) return img1, img2 + class CoupledDiscriminators(nn.Module): def __init__(self): super(CoupledDiscriminators, self).__init__() @@ -102,9 +105,9 @@ def discriminator_block(in_filters, out_filters, bn=True): *discriminator_block(64, 128), ) # The height and width of downsampled image - ds_size = opt.img_size // 2**4 - self.D1 = nn.Linear(128*ds_size**2, 1) - self.D2 = nn.Linear(128*ds_size**2, 1) + ds_size = opt.img_size // 2 ** 4 + self.D1 = nn.Linear(128 * ds_size ** 2, 1) + self.D2 = nn.Linear(128 * ds_size ** 2, 1) def forward(self, img1, img2): # Determine validity of first image @@ -118,6 +121,7 @@ def forward(self, img1, img2): return validity1, validity2 + # Loss function adversarial_loss = torch.nn.MSELoss() @@ -134,25 +138,37 @@ def forward(self, img1, img2): coupled_discriminators.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader1 = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) - -os.makedirs('../../data/mnistm', exist_ok=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) + +os.makedirs("../../data/mnistm", exist_ok=True) dataloader2 = torch.utils.data.DataLoader( - mnistm.MNISTM('../../data/mnistm', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + mnistm.MNISTM( + "../../data/mnistm", + train=True, + download=True, + transform=transforms.Compose( + [ + transforms.Resize(opt.img_size), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), + ] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(coupled_generators.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -191,8 +207,7 @@ def forward(self, img1, img2): # Determine validity of generated images validity1, validity2 = coupled_discriminators(gen_imgs1, gen_imgs2) - g_loss = (adversarial_loss(validity1, valid) + \ - adversarial_loss(validity2, valid)) / 2 + g_loss = (adversarial_loss(validity1, valid) + adversarial_loss(validity2, valid)) / 2 g_loss.backward() optimizer_G.step() @@ -207,18 +222,22 @@ def forward(self, img1, img2): validity1_real, validity2_real = coupled_discriminators(imgs1, imgs2) validity1_fake, validity2_fake = coupled_discriminators(gen_imgs1.detach(), gen_imgs2.detach()) - d_loss = (adversarial_loss(validity1_real, valid) + \ - adversarial_loss(validity1_fake, fake) + \ - adversarial_loss(validity2_real, valid) + \ - adversarial_loss(validity2_fake, fake)) / 4 + d_loss = ( + adversarial_loss(validity1_real, valid) + + adversarial_loss(validity1_fake, fake) + + adversarial_loss(validity2_real, valid) + + adversarial_loss(validity2_fake, fake) + ) / 4 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(dataloader1), - 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(dataloader1), d_loss.item(), g_loss.item()) + ) batches_done = epoch * len(dataloader1) + i if batches_done % opt.sample_interval == 0: gen_imgs = torch.cat((gen_imgs1.data, gen_imgs2.data), 0) - save_image(gen_imgs, 'images/%d.png' % batches_done, nrow=8, normalize=True) + save_image(gen_imgs, "images/%d.png" % batches_done, nrow=8, normalize=True) diff --git a/implementations/cogan/mnistm.py b/implementations/cogan/mnistm.py index b55b29d9..e0a9832f 100644 --- a/implementations/cogan/mnistm.py +++ b/implementations/cogan/mnistm.py @@ -21,16 +21,12 @@ class MNISTM(data.Dataset): url = "https://github.com/VanushVaswani/keras_mnistm/releases/download/1.0/keras_mnistm.pkl.gz" - raw_folder = 'raw' - processed_folder = 'processed' - training_file = 'mnist_m_train.pt' - test_file = 'mnist_m_test.pt' - - def __init__(self, - root, mnist_root="data", - train=True, - transform=None, target_transform=None, - download=False): + raw_folder = "raw" + processed_folder = "processed" + training_file = "mnist_m_train.pt" + test_file = "mnist_m_test.pt" + + def __init__(self, root, mnist_root="data", train=True, transform=None, target_transform=None, download=False): """Init MNIST-M dataset.""" super(MNISTM, self).__init__() self.root = os.path.expanduser(root) @@ -43,19 +39,16 @@ def __init__(self, self.download() if not self._check_exists(): - raise RuntimeError('Dataset not found.' + - ' You can use download=True to download it') + raise RuntimeError("Dataset not found." + " You can use download=True to download it") if self.train: - self.train_data, self.train_labels = \ - torch.load(os.path.join(self.root, - self.processed_folder, - self.training_file)) + self.train_data, self.train_labels = torch.load( + os.path.join(self.root, self.processed_folder, self.training_file) + ) else: - self.test_data, self.test_labels = \ - torch.load(os.path.join(self.root, - self.processed_folder, - self.test_file)) + self.test_data, self.test_labels = torch.load( + os.path.join(self.root, self.processed_folder, self.test_file) + ) def __getitem__(self, index): """Get images and target for data loader. @@ -73,7 +66,7 @@ def __getitem__(self, index): # doing this so that it is consistent with all other datasets # to return a PIL Image - img = Image.fromarray(img.squeeze().numpy(), mode='RGB') + img = Image.fromarray(img.squeeze().numpy(), mode="RGB") if self.transform is not None: img = self.transform(img) @@ -91,12 +84,9 @@ def __len__(self): return len(self.test_data) def _check_exists(self): - return os.path.exists(os.path.join(self.root, - self.processed_folder, - self.training_file)) and \ - os.path.exists(os.path.join(self.root, - self.processed_folder, - self.test_file)) + return os.path.exists(os.path.join(self.root, self.processed_folder, self.training_file)) and os.path.exists( + os.path.join(self.root, self.processed_folder, self.test_file) + ) def download(self): """Download the MNIST data.""" @@ -121,45 +111,36 @@ def download(self): raise # download pkl files - print('Downloading ' + self.url) - filename = self.url.rpartition('/')[2] + print("Downloading " + self.url) + filename = self.url.rpartition("/")[2] file_path = os.path.join(self.root, self.raw_folder, filename) - if not os.path.exists(file_path.replace('.gz', '')): + if not os.path.exists(file_path.replace(".gz", "")): data = urllib.request.urlopen(self.url) - with open(file_path, 'wb') as f: + with open(file_path, "wb") as f: f.write(data.read()) - with open(file_path.replace('.gz', ''), 'wb') as out_f, \ - gzip.GzipFile(file_path) as zip_f: + with open(file_path.replace(".gz", ""), "wb") as out_f, gzip.GzipFile(file_path) as zip_f: out_f.write(zip_f.read()) os.unlink(file_path) # process and save as torch files - print('Processing...') + print("Processing...") # load MNIST-M images from pkl file - with open(file_path.replace('.gz', ''), "rb") as f: - mnist_m_data = pickle.load(f, encoding='bytes') - mnist_m_train_data = torch.ByteTensor(mnist_m_data[b'train']) - mnist_m_test_data = torch.ByteTensor(mnist_m_data[b'test']) + with open(file_path.replace(".gz", ""), "rb") as f: + mnist_m_data = pickle.load(f, encoding="bytes") + mnist_m_train_data = torch.ByteTensor(mnist_m_data[b"train"]) + mnist_m_test_data = torch.ByteTensor(mnist_m_data[b"test"]) # get MNIST labels - mnist_train_labels = datasets.MNIST(root=self.mnist_root, - train=True, - download=True).train_labels - mnist_test_labels = datasets.MNIST(root=self.mnist_root, - train=False, - download=True).test_labels + mnist_train_labels = datasets.MNIST(root=self.mnist_root, train=True, download=True).train_labels + mnist_test_labels = datasets.MNIST(root=self.mnist_root, train=False, download=True).test_labels # save MNIST-M dataset training_set = (mnist_m_train_data, mnist_train_labels) test_set = (mnist_m_test_data, mnist_test_labels) - with open(os.path.join(self.root, - self.processed_folder, - self.training_file), 'wb') as f: + with open(os.path.join(self.root, self.processed_folder, self.training_file), "wb") as f: torch.save(training_set, f) - with open(os.path.join(self.root, - self.processed_folder, - self.test_file), 'wb') as f: + with open(os.path.join(self.root, self.processed_folder, self.test_file), "wb") as f: torch.save(test_set, f) - print('Done!') + print("Done!") diff --git a/implementations/dcgan/dcgan.py b/implementations/dcgan/dcgan.py index 2e772d5e..46e64c06 100644 --- a/implementations/dcgan/dcgan.py +++ b/implementations/dcgan/dcgan.py @@ -14,38 +14,40 @@ 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=64, help='size of the batches') -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=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, 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=64, help="size of the batches") +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=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling") opt = parser.parse_args() print(opt) cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm2d') != -1: + 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) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.init_size = opt.img_size // 4 - self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.BatchNorm2d(128), @@ -58,7 +60,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, z): @@ -67,14 +69,13 @@ def forward(self, z): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() def discriminator_block(in_filters, out_filters, bn=True): - block = [ nn.Conv2d(in_filters, out_filters, 3, 2, 1), - nn.LeakyReLU(0.2, inplace=True), - nn.Dropout2d(0.25)] + block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)] if bn: block.append(nn.BatchNorm2d(out_filters, 0.8)) return block @@ -87,9 +88,8 @@ def discriminator_block(in_filters, out_filters, bn=True): ) # The height and width of downsampled image - ds_size = opt.img_size // 2**4 - self.adv_layer = nn.Sequential( nn.Linear(128*ds_size**2, 1), - nn.Sigmoid()) + ds_size = opt.img_size // 2 ** 4 + self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid()) def forward(self, img): out = self.model(img) @@ -98,6 +98,7 @@ def forward(self, img): return validity + # Loss function adversarial_loss = torch.nn.BCELoss() @@ -115,15 +116,19 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -177,9 +182,11 @@ def forward(self, img): 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()) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/dragan/dragan.py b/implementations/dragan/dragan.py index 1ef85b67..6bae1631 100644 --- a/implementations/dragan/dragan.py +++ b/implementations/dragan/dragan.py @@ -15,38 +15,40 @@ 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=64, help='size of the batches') -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=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=1000, 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=64, help="size of the batches") +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=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=1000, help="interval between image sampling") opt = parser.parse_args() print(opt) cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm2d') != -1: + 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) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.init_size = opt.img_size // 4 - self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.BatchNorm2d(128), @@ -59,7 +61,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, noise): @@ -68,14 +70,13 @@ def forward(self, noise): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() def discriminator_block(in_filters, out_filters, bn=True): - block = [ nn.Conv2d(in_filters, out_filters, 3, 2, 1), - nn.LeakyReLU(0.2, inplace=True), - nn.Dropout2d(0.25)] + block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)] if bn: block.append(nn.BatchNorm2d(out_filters, 0.8)) return block @@ -88,9 +89,8 @@ def discriminator_block(in_filters, out_filters, bn=True): ) # The height and width of downsampled image - ds_size = opt.img_size // 2**4 - self.adv_layer = nn.Sequential( nn.Linear(128*ds_size**2, 1), - nn.Sigmoid()) + ds_size = opt.img_size // 2 ** 4 + self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid()) def forward(self, img): out = self.model(img) @@ -99,6 +99,7 @@ def forward(self, img): return validity + # Loss function adversarial_loss = torch.nn.BCELoss() @@ -119,15 +120,19 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) -mnist_loader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) +os.makedirs("../../data/mnist", exist_ok=True) +dataloader = torch.utils.data.DataLoader( + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -135,12 +140,13 @@ def forward(self, img): Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor + def compute_gradient_penalty(D, X): """Calculates the gradient penalty loss for DRAGAN""" # Random weight term for interpolation alpha = Tensor(np.random.random(size=X.shape)) - interpolates = alpha * X + ((1 - alpha) * (X + 0.5 * X.std() * torch.rand(X.size()))) + interpolates = alpha * X + ((1 - alpha) * (X + 0.5 * X.std() * torch.rand(X.size()))) interpolates = Variable(interpolates, requires_grad=True) d_interpolates = D(interpolates) @@ -148,13 +154,19 @@ def compute_gradient_penalty(D, X): fake = Variable(Tensor(X.shape[0], 1).fill_(1.0), requires_grad=False) # Get gradient w.r.t. interpolates - gradients = autograd.grad(outputs=d_interpolates, inputs=interpolates, - grad_outputs=fake, create_graph=True, retain_graph=True, - only_inputs=True)[0] + gradients = autograd.grad( + outputs=d_interpolates, + inputs=interpolates, + grad_outputs=fake, + create_graph=True, + retain_graph=True, + only_inputs=True, + )[0] gradient_penalty = lambda_gp * ((gradients.norm(2, dim=1) - 1) ** 2).mean() return gradient_penalty + # ---------- # Training # ---------- @@ -204,7 +216,9 @@ def compute_gradient_penalty(D, X): optimizer_D.step() - print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" % (epoch, opt.n_epochs, i, len(mnist_loader), - 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(mnist_loader), d_loss.item(), g_loss.item()) + ) - save_image(gen_imgs.data, 'images/%d.png' % epoch, nrow=int(math.sqrt(opt.batch_size)), normalize=True) + save_image(gen_imgs.data, "images/%d.png" % epoch, nrow=int(math.sqrt(opt.batch_size)), normalize=True) diff --git a/implementations/ebgan/ebgan.py b/implementations/ebgan/ebgan.py index 72891885..0392930b 100644 --- a/implementations/ebgan/ebgan.py +++ b/implementations/ebgan/ebgan.py @@ -14,19 +14,19 @@ 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=64, help='size of the batches') -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=62, help='dimensionality of the latent space') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, help='number of image channels') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +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=62, help="dimensionality of the latent space") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="number of image channels") opt = parser.parse_args() print(opt) @@ -34,20 +34,22 @@ cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm2d') != -1: + 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) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.init_size = opt.img_size // 4 - self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.Upsample(scale_factor=2), @@ -59,7 +61,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, noise): @@ -68,18 +70,16 @@ def forward(self, noise): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() # Upsampling - self.down = nn.Sequential( - nn.Conv2d(opt.channels, 64, 3, 2, 1), - nn.ReLU(), - ) + self.down = nn.Sequential(nn.Conv2d(opt.channels, 64, 3, 2, 1), nn.ReLU()) # Fully-connected layers - self.down_size = (opt.img_size // 2) - down_dim = 64 * (opt.img_size // 2)**2 + self.down_size = opt.img_size // 2 + down_dim = 64 * (opt.img_size // 2) ** 2 self.embedding = nn.Linear(down_dim, 32) @@ -88,13 +88,10 @@ def __init__(self): nn.ReLU(inplace=True), nn.Linear(32, down_dim), nn.BatchNorm1d(down_dim), - nn.ReLU(inplace=True) + nn.ReLU(inplace=True), ) # Upsampling - self.up = nn.Sequential( - nn.Upsample(scale_factor=2), - nn.Conv2d(64, opt.channels, 3, 1, 1) - ) + self.up = nn.Sequential(nn.Upsample(scale_factor=2), nn.Conv2d(64, opt.channels, 3, 1, 1)) def forward(self, img): out = self.down(img) @@ -103,6 +100,7 @@ def forward(self, img): out = self.up(out.view(out.size(0), 64, self.down_size, self.down_size)) return out, embedding + # Reconstruction loss of AE pixelwise_loss = nn.MSELoss() @@ -120,15 +118,19 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -136,6 +138,7 @@ def forward(self, img): Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor + def pullaway_loss(embeddings): norm = torch.sqrt(torch.sum(embeddings ** 2, -1, keepdim=True)) normalized_emb = embeddings / norm @@ -144,13 +147,14 @@ def pullaway_loss(embeddings): loss_pt = (torch.sum(similarity) - batch_size) / (batch_size * (batch_size - 1)) return loss_pt + # ---------- # Training # ---------- # BEGAN hyper parameters lambda_pt = 0.1 -margin = max(1, opt.batch_size / 64.) +margin = max(1, opt.batch_size / 64.0) for epoch in range(opt.n_epochs): for i, (imgs, _) in enumerate(dataloader): @@ -172,8 +176,7 @@ def pullaway_loss(embeddings): recon_imgs, img_embeddings = discriminator(gen_imgs) # Loss measures generator's ability to fool the discriminator - g_loss = pixelwise_loss(recon_imgs, gen_imgs.detach()) + \ - lambda_pt * pullaway_loss(img_embeddings) + g_loss = pixelwise_loss(recon_imgs, gen_imgs.detach()) + lambda_pt * pullaway_loss(img_embeddings) g_loss.backward() optimizer_G.step() @@ -198,13 +201,15 @@ def pullaway_loss(embeddings): d_loss.backward() optimizer_D.step() - #-------------- + # -------------- # Log Progress - #-------------- + # -------------- - 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()) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/gan/gan.py b/implementations/gan/gan.py index 1b4a3887..d6f1d935 100644 --- a/implementations/gan/gan.py +++ b/implementations/gan/gan.py @@ -14,19 +14,19 @@ 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=64, help='size of the batches') -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=28, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, help='interval betwen image samples') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +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=28, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples") opt = parser.parse_args() print(opt) @@ -34,6 +34,7 @@ cuda = True if torch.cuda.is_available() else False + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() @@ -59,6 +60,7 @@ def forward(self, z): img = img.view(img.size(0), *img_shape) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() @@ -69,7 +71,7 @@ def __init__(self): nn.Linear(512, 256), nn.LeakyReLU(0.2, inplace=True), nn.Linear(256, 1), - nn.Sigmoid() + nn.Sigmoid(), ) def forward(self, img): @@ -78,6 +80,7 @@ def forward(self, img): return validity + # Loss function adversarial_loss = torch.nn.BCELoss() @@ -91,14 +94,19 @@ def forward(self, img): adversarial_loss.cuda() # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -152,9 +160,11 @@ def forward(self, img): 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()) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/infogan/infogan.py b/implementations/infogan/infogan.py index 45190be2..38a2f530 100644 --- a/implementations/infogan/infogan.py +++ b/implementations/infogan/infogan.py @@ -15,50 +15,53 @@ import torch.nn.functional as F import torch -os.makedirs('images/static/', exist_ok=True) -os.makedirs('images/varying_c1/', exist_ok=True) -os.makedirs('images/varying_c2/', exist_ok=True) +os.makedirs("images/static/", exist_ok=True) +os.makedirs("images/varying_c1/", exist_ok=True) +os.makedirs("images/varying_c2/", 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=64, help='size of the batches') -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=62, help='dimensionality of the latent space') -parser.add_argument('--code_dim', type=int, default=2, help='latent code') -parser.add_argument('--n_classes', type=int, default=10, help='number of classes for dataset') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, 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=64, help="size of the batches") +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=62, help="dimensionality of the latent space") +parser.add_argument("--code_dim", type=int, default=2, help="latent code") +parser.add_argument("--n_classes", type=int, default=10, help="number of classes for dataset") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling") opt = parser.parse_args() print(opt) cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm') != -1: + elif classname.find("BatchNorm") != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0) + def to_categorical(y, num_columns): """Returns one-hot encoded Variable""" y_cat = np.zeros((y.shape[0], num_columns)) - y_cat[range(y.shape[0]), y] = 1. + y_cat[range(y.shape[0]), y] = 1.0 return Variable(FloatTensor(y_cat)) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() input_dim = opt.latent_dim + opt.n_classes + opt.code_dim - self.init_size = opt.img_size // 4 # Initial size before upsampling - self.l1 = nn.Sequential(nn.Linear(input_dim, 128*self.init_size**2)) + self.init_size = opt.img_size // 4 # Initial size before upsampling + self.l1 = nn.Sequential(nn.Linear(input_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.BatchNorm2d(128), @@ -71,7 +74,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, noise, labels, code): @@ -81,15 +84,14 @@ def forward(self, noise, labels, code): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() def discriminator_block(in_filters, out_filters, bn=True): """Returns layers of each discriminator block""" - block = [ nn.Conv2d(in_filters, out_filters, 3, 2, 1), - nn.LeakyReLU(0.2, inplace=True), - nn.Dropout2d(0.25)] + block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)] if bn: block.append(nn.BatchNorm2d(out_filters, 0.8)) return block @@ -102,15 +104,12 @@ def discriminator_block(in_filters, out_filters, bn=True): ) # The height and width of downsampled image - ds_size = opt.img_size // 2**4 + ds_size = opt.img_size // 2 ** 4 # Output layers - self.adv_layer = nn.Sequential(nn.Linear(128*ds_size**2, 1)) - self.aux_layer = nn.Sequential( - nn.Linear(128*ds_size**2, opt.n_classes), - nn.Softmax() - ) - self.latent_layer = nn.Sequential(nn.Linear(128*ds_size**2, opt.code_dim)) + self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1)) + self.aux_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.n_classes), nn.Softmax()) + self.latent_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.code_dim)) def forward(self, img): out = self.conv_blocks(img) @@ -121,6 +120,7 @@ def forward(self, img): return validity, label, latent_code + # Loss functions adversarial_loss = torch.nn.MSELoss() categorical_loss = torch.nn.CrossEntropyLoss() @@ -146,47 +146,55 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # 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)) -optimizer_info = torch.optim.Adam(itertools.chain(generator.parameters(), discriminator.parameters()), - lr=opt.lr, betas=(opt.b1, opt.b2)) +optimizer_info = torch.optim.Adam( + itertools.chain(generator.parameters(), discriminator.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2) +) FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor # Static generator inputs for sampling -static_z = Variable(FloatTensor(np.zeros((opt.n_classes**2, opt.latent_dim)))) -static_label = to_categorical(np.array([num for _ in range(opt.n_classes) for num in range(opt.n_classes)]), - num_columns=opt.n_classes) -static_code = Variable(FloatTensor(np.zeros((opt.n_classes**2, opt.code_dim)))) +static_z = Variable(FloatTensor(np.zeros((opt.n_classes ** 2, opt.latent_dim)))) +static_label = to_categorical( + np.array([num for _ in range(opt.n_classes) for num in range(opt.n_classes)]), num_columns=opt.n_classes +) +static_code = Variable(FloatTensor(np.zeros((opt.n_classes ** 2, opt.code_dim)))) + def sample_image(n_row, batches_done): """Saves a grid of generated digits ranging from 0 to n_classes""" # Static sample - z = Variable(FloatTensor(np.random.normal(0, 1, (n_row**2, opt.latent_dim)))) + z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim)))) static_sample = generator(z, static_label, static_code) - save_image(static_sample.data, 'images/static/%d.png' % batches_done, nrow=n_row, normalize=True) + save_image(static_sample.data, "images/static/%d.png" % batches_done, nrow=n_row, normalize=True) # Get varied c1 and c2 - zeros = np.zeros((n_row**2, 1)) + zeros = np.zeros((n_row ** 2, 1)) c_varied = np.repeat(np.linspace(-1, 1, n_row)[:, np.newaxis], n_row, 0) c1 = Variable(FloatTensor(np.concatenate((c_varied, zeros), -1))) c2 = Variable(FloatTensor(np.concatenate((zeros, c_varied), -1))) sample1 = generator(static_z, static_label, c1) sample2 = generator(static_z, static_label, c2) - save_image(sample1.data, 'images/varying_c1/%d.png' % batches_done, nrow=n_row, normalize=True) - save_image(sample2.data, 'images/varying_c2/%d.png' % batches_done, nrow=n_row, normalize=True) + save_image(sample1.data, "images/varying_c1/%d.png" % batches_done, nrow=n_row, normalize=True) + save_image(sample2.data, "images/varying_c2/%d.png" % batches_done, nrow=n_row, normalize=True) + # ---------- # Training @@ -246,9 +254,9 @@ def sample_image(n_row, batches_done): d_loss.backward() optimizer_D.step() - #------------------ + # ------------------ # Information Loss - #------------------ + # ------------------ optimizer_info.zero_grad() @@ -258,7 +266,6 @@ def sample_image(n_row, batches_done): # Ground truth labels gt_labels = Variable(LongTensor(sampled_labels), requires_grad=False) - # Sample noise, labels and code as generator input z = Variable(FloatTensor(np.random.normal(0, 1, (batch_size, opt.latent_dim)))) label_input = to_categorical(sampled_labels, num_columns=opt.n_classes) @@ -267,18 +274,21 @@ def sample_image(n_row, batches_done): gen_imgs = generator(z, label_input, code_input) _, pred_label, pred_code = discriminator(gen_imgs) - info_loss = lambda_cat * categorical_loss(pred_label, gt_labels) + \ - lambda_con * continuous_loss(pred_code, code_input) + info_loss = lambda_cat * categorical_loss(pred_label, gt_labels) + lambda_con * continuous_loss( + pred_code, code_input + ) info_loss.backward() optimizer_info.step() - #-------------- + # -------------- # Log Progress - #-------------- + # -------------- - print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] [info loss: %f]" % (epoch, opt.n_epochs, i, len(dataloader), - d_loss.item(), g_loss.item(), info_loss.item())) + print( + "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] [info loss: %f]" + % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item(), info_loss.item()) + ) batches_done = epoch * len(dataloader) + i if batches_done % opt.sample_interval == 0: sample_image(n_row=10, batches_done=batches_done) diff --git a/implementations/lsgan/lsgan.py b/implementations/lsgan/lsgan.py index 7f882aab..1930ab84 100644 --- a/implementations/lsgan/lsgan.py +++ b/implementations/lsgan/lsgan.py @@ -14,38 +14,40 @@ 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=64, help='size of the batches') -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=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=1000, help='number of image channels') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +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=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=1000, help="number of image channels") opt = parser.parse_args() print(opt) cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm') != -1: + elif classname.find("BatchNorm") != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.init_size = opt.img_size // 4 - self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.Upsample(scale_factor=2), @@ -57,7 +59,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, z): @@ -66,14 +68,13 @@ def forward(self, z): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() def discriminator_block(in_filters, out_filters, bn=True): - block = [ nn.Conv2d(in_filters, out_filters, 3, 2, 1), - nn.LeakyReLU(0.2, inplace=True), - nn.Dropout2d(0.25)] + block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)] if bn: block.append(nn.BatchNorm2d(out_filters, 0.8)) return block @@ -86,8 +87,8 @@ def discriminator_block(in_filters, out_filters, bn=True): ) # The height and width of downsampled image - ds_size = opt.img_size // 2**4 - self.adv_layer = nn.Linear(128*ds_size**2, 1) + ds_size = opt.img_size // 2 ** 4 + self.adv_layer = nn.Linear(128 * ds_size ** 2, 1) def forward(self, img): out = self.model(img) @@ -96,6 +97,7 @@ def forward(self, img): return validity + # !!! Minimizes MSE instead of BCE adversarial_loss = torch.nn.MSELoss() @@ -113,15 +115,19 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -175,9 +181,11 @@ def forward(self, img): 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()) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/pixelda/pixelda.py b/implementations/pixelda/pixelda.py index 3e18eaee..eb80b2d0 100644 --- a/implementations/pixelda/pixelda.py +++ b/implementations/pixelda/pixelda.py @@ -17,38 +17,40 @@ 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=64, help='size of the batches') -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('--n_residual_blocks', type=int, default=6, help='number of residual blocks in generator') -parser.add_argument('--latent_dim', type=int, default=10, help='dimensionality of the noise input') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=3, help='number of image channels') -parser.add_argument('--n_classes', type=int, default=10, help='number of classes in the dataset') -parser.add_argument('--sample_interval', type=int, default=300, help='interval betwen image samples') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +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("--n_residual_blocks", type=int, default=6, help="number of residual blocks in generator") +parser.add_argument("--latent_dim", type=int, default=10, help="dimensionality of the noise input") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=3, help="number of image channels") +parser.add_argument("--n_classes", type=int, default=10, help="number of classes in the dataset") +parser.add_argument("--sample_interval", type=int, default=300, help="interval betwen image samples") opt = parser.parse_args() print(opt) # Calculate output of image discriminator (PatchGAN) -patch = int(opt.img_size / 2**4) +patch = int(opt.img_size / 2 ** 4) patch = (1, patch, patch) cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm') != -1: + elif classname.find("BatchNorm") != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0) + class ResidualBlock(nn.Module): def __init__(self, in_features=64, out_features=64): super(ResidualBlock, self).__init__() @@ -58,20 +60,21 @@ def __init__(self, in_features=64, out_features=64): nn.BatchNorm2d(in_features), nn.ReLU(inplace=True), nn.Conv2d(in_features, in_features, 3, 1, 1), - nn.BatchNorm2d(in_features) + nn.BatchNorm2d(in_features), ) def forward(self, x): return x + self.block(x) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() # Fully-connected layer which constructs image channel shaped output from noise - self.fc = nn.Linear(opt.latent_dim, opt.channels*opt.img_size**2) + self.fc = nn.Linear(opt.latent_dim, opt.channels * opt.img_size ** 2) - self.l1 = nn.Sequential(nn.Conv2d(opt.channels*2, 64, 3, 1, 1), nn.ReLU(inplace=True)) + 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): @@ -80,7 +83,6 @@ def __init__(self): self.l2 = nn.Sequential(nn.Conv2d(64, opt.channels, 3, 1, 1), nn.Tanh()) - def forward(self, img, z): gen_input = torch.cat((img, self.fc(z).view(*img.shape)), 1) out = self.l1(gen_input) @@ -89,14 +91,14 @@ def forward(self, img, z): return img_ + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() def block(in_features, out_features, normalization=True): """Discriminator block""" - layers = [ nn.Conv2d(in_features, out_features, 3, stride=2, padding=1), - nn.LeakyReLU(0.2, inplace=True) ] + layers = [nn.Conv2d(in_features, out_features, 3, stride=2, padding=1), nn.LeakyReLU(0.2, inplace=True)] if normalization: layers.append(nn.InstanceNorm2d(out_features)) return layers @@ -121,24 +123,17 @@ def __init__(self): def block(in_features, out_features, normalization=True): """Classifier block""" - layers = [ nn.Conv2d(in_features, out_features, 3, stride=2, padding=1), - nn.LeakyReLU(0.2, inplace=True) ] + layers = [nn.Conv2d(in_features, out_features, 3, stride=2, padding=1), nn.LeakyReLU(0.2, inplace=True)] if normalization: layers.append(nn.InstanceNorm2d(out_features)) return layers self.model = nn.Sequential( - *block(opt.channels, 64, normalization=False), - *block(64, 128), - *block(128, 256), - *block(256, 512) + *block(opt.channels, 64, normalization=False), *block(64, 128), *block(128, 256), *block(256, 512) ) - input_size = opt.img_size // 2**4 - self.output_layer = nn.Sequential( - nn.Linear(512*input_size**2, opt.n_classes), - nn.Softmax() - ) + input_size = opt.img_size // 2 ** 4 + self.output_layer = nn.Sequential(nn.Linear(512 * input_size ** 2, opt.n_classes), nn.Softmax()) def forward(self, img): feature_repr = self.model(img) @@ -146,12 +141,13 @@ def forward(self, img): label = self.output_layer(feature_repr) return label + # Loss function adversarial_loss = torch.nn.MSELoss() task_loss = torch.nn.CrossEntropyLoss() # Loss weights -lambda_adv = 1 +lambda_adv = 1 lambda_task = 0.1 # Initialize generator and discriminator @@ -172,29 +168,42 @@ def forward(self, img): classifier.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader_A = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) - -os.makedirs('../../data/mnistm', exist_ok=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) + +os.makedirs("../../data/mnistm", exist_ok=True) dataloader_B = torch.utils.data.DataLoader( - MNISTM('../../data/mnistm', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + MNISTM( + "../../data/mnistm", + train=True, + download=True, + transform=transforms.Compose( + [ + transforms.Resize(opt.img_size), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), + ] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers -optimizer_G = torch.optim.Adam( itertools.chain(generator.parameters(), classifier.parameters()), - lr=opt.lr, betas=(opt.b1, opt.b2)) +optimizer_G = torch.optim.Adam( + itertools.chain(generator.parameters(), classifier.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2) +) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor @@ -218,9 +227,9 @@ def forward(self, img): fake = Variable(FloatTensor(batch_size, *patch).fill_(0.0), requires_grad=False) # Configure input - imgs_A = Variable(imgs_A.type(FloatTensor).expand(batch_size, 3, opt.img_size, opt.img_size)) - labels_A = Variable(labels_A.type(LongTensor)) - imgs_B = Variable(imgs_B.type(FloatTensor)) + imgs_A = Variable(imgs_A.type(FloatTensor).expand(batch_size, 3, opt.img_size, opt.img_size)) + labels_A = Variable(labels_A.type(LongTensor)) + imgs_B = Variable(imgs_B.type(FloatTensor)) # ----------------- # Train Generator @@ -238,12 +247,10 @@ def forward(self, img): label_pred = classifier(fake_B) # Calculate the task loss - task_loss_ = (task_loss(label_pred, labels_A) + \ - task_loss(classifier(imgs_A), labels_A)) / 2 + task_loss_ = (task_loss(label_pred, labels_A) + task_loss(classifier(imgs_A), labels_A)) / 2 # Loss measures generator's ability to fool the discriminator - g_loss = lambda_adv * adversarial_loss(discriminator(fake_B), valid) + \ - lambda_task * task_loss_ + g_loss = lambda_adv * adversarial_loss(discriminator(fake_B), valid) + lambda_task * task_loss_ g_loss.backward() optimizer_G.step() @@ -279,14 +286,23 @@ def forward(self, img): if len(target_performance) > 100: target_performance.pop(0) - print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] [CLF acc: %3d%% (%3d%%), target_acc: %3d%% (%3d%%)]" % - (epoch, opt.n_epochs, - i, len(dataloader_A), - d_loss.item(), g_loss.item(), - 100*acc, 100*np.mean(task_performance), - 100*target_acc, 100*np.mean(target_performance))) + print( + "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] [CLF acc: %3d%% (%3d%%), target_acc: %3d%% (%3d%%)]" + % ( + epoch, + opt.n_epochs, + i, + len(dataloader_A), + d_loss.item(), + g_loss.item(), + 100 * acc, + 100 * np.mean(task_performance), + 100 * target_acc, + 100 * np.mean(target_performance), + ) + ) batches_done = len(dataloader_A) * epoch + i if batches_done % opt.sample_interval == 0: sample = torch.cat((imgs_A.data[:5], fake_B.data[:5], imgs_B.data[:5]), -2) - save_image(sample, 'images/%d.png' % batches_done, nrow=int(math.sqrt(batch_size)), normalize=True) + save_image(sample, "images/%d.png" % batches_done, nrow=int(math.sqrt(batch_size)), normalize=True) diff --git a/implementations/sgan/sgan.py b/implementations/sgan/sgan.py index ba6bed84..d12bf292 100644 --- a/implementations/sgan/sgan.py +++ b/implementations/sgan/sgan.py @@ -14,41 +14,43 @@ 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=64, help='size of the batches') -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('--num_classes', type=int, default=10, help='number of classes for dataset') -parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, 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=64, help="size of the batches") +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("--num_classes", type=int, default=10, help="number of classes for dataset") +parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling") opt = parser.parse_args() print(opt) cuda = True if torch.cuda.is_available() else False + def weights_init_normal(m): classname = m.__class__.__name__ - if classname.find('Conv') != -1: + if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) - elif classname.find('BatchNorm') != -1: + elif classname.find("BatchNorm") != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0) + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.label_emb = nn.Embedding(opt.num_classes, opt.latent_dim) - self.init_size = opt.img_size // 4 # Initial size before upsampling - self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2)) + self.init_size = opt.img_size // 4 # Initial size before upsampling + self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2)) self.conv_blocks = nn.Sequential( nn.BatchNorm2d(128), @@ -61,7 +63,7 @@ def __init__(self): nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, opt.channels, 3, stride=1, padding=1), - nn.Tanh() + nn.Tanh(), ) def forward(self, noise): @@ -70,15 +72,14 @@ def forward(self, noise): img = self.conv_blocks(out) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() def discriminator_block(in_filters, out_filters, bn=True): """Returns layers of each discriminator block""" - block = [ nn.Conv2d(in_filters, out_filters, 3, 2, 1), - nn.LeakyReLU(0.2, inplace=True), - nn.Dropout2d(0.25)] + block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)] if bn: block.append(nn.BatchNorm2d(out_filters, 0.8)) return block @@ -91,13 +92,11 @@ def discriminator_block(in_filters, out_filters, bn=True): ) # The height and width of downsampled image - ds_size = opt.img_size // 2**4 + ds_size = opt.img_size // 2 ** 4 # Output layers - self.adv_layer = nn.Sequential( nn.Linear(128*ds_size**2, 1), - nn.Sigmoid()) - self.aux_layer = nn.Sequential( nn.Linear(128*ds_size**2, opt.num_classes+1), - nn.Softmax()) + self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid()) + self.aux_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.num_classes + 1), nn.Softmax()) def forward(self, img): out = self.conv_blocks(img) @@ -107,6 +106,7 @@ def forward(self, img): return validity, label + # Loss functions adversarial_loss = torch.nn.BCELoss() auxiliary_loss = torch.nn.CrossEntropyLoss() @@ -126,15 +126,19 @@ def forward(self, img): discriminator.apply(weights_init_normal) # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.Resize(opt.img_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -188,13 +192,11 @@ def forward(self, img): # Loss for real images real_pred, real_aux = discriminator(real_imgs) - d_real_loss = (adversarial_loss(real_pred, valid) + \ - auxiliary_loss(real_aux, labels)) / 2 + d_real_loss = (adversarial_loss(real_pred, valid) + auxiliary_loss(real_aux, labels)) / 2 # Loss for fake images fake_pred, fake_aux = discriminator(gen_imgs.detach()) - d_fake_loss = (adversarial_loss(fake_pred, fake) + \ - auxiliary_loss(fake_aux, fake_aux_gt)) / 2 + d_fake_loss = (adversarial_loss(fake_pred, fake) + auxiliary_loss(fake_aux, fake_aux_gt)) / 2 # Total discriminator loss d_loss = (d_real_loss + d_fake_loss) / 2 @@ -207,10 +209,11 @@ def forward(self, img): d_loss.backward() optimizer_D.step() - print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]" % (epoch, opt.n_epochs, i, len(dataloader), - d_loss.item(), 100 * d_acc, - g_loss.item())) + print( + "[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]" + % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), 100 * d_acc, g_loss.item()) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/softmax_gan/softmax_gan.py b/implementations/softmax_gan/softmax_gan.py index b2c26e64..23a0775b 100644 --- a/implementations/softmax_gan/softmax_gan.py +++ b/implementations/softmax_gan/softmax_gan.py @@ -14,19 +14,19 @@ 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=64, help='size of the batches') -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=28, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--sample_interval', type=int, default=400, help='interval betwen image samples') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +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=28, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples") opt = parser.parse_args() print(opt) @@ -34,12 +34,13 @@ cuda = True if torch.cuda.is_available() else False + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() def block(in_feat, out_feat, normalize=True): - layers = [ nn.Linear(in_feat, out_feat)] + layers = [nn.Linear(in_feat, out_feat)] if normalize: layers.append(nn.BatchNorm1d(out_feat, 0.8)) layers.append(nn.LeakyReLU(0.2, inplace=True)) @@ -59,16 +60,17 @@ def forward(self, z): img = img.view(img.shape[0], *img_shape) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() self.model = nn.Sequential( - nn.Linear(opt.img_size**2, 512), + nn.Linear(opt.img_size ** 2, 512), nn.LeakyReLU(0.2, inplace=True), nn.Linear(512, 256), nn.LeakyReLU(0.2, inplace=True), - nn.Linear(256, 1) + nn.Linear(256, 1), ) def forward(self, img): @@ -77,6 +79,7 @@ def forward(self, img): return validity + # Loss function adversarial_loss = torch.nn.BCELoss() @@ -90,14 +93,19 @@ def forward(self, img): adversarial_loss.cuda() # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) @@ -105,9 +113,11 @@ def forward(self, img): Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor + def log(x): return torch.log(x + 1e-8) + # ---------- # Training # ---------- @@ -121,7 +131,7 @@ def log(x): batch_size = imgs.shape[0] # Adversarial ground truths - g_target = 1 / (batch_size*2) + g_target = 1 / (batch_size * 2) d_target = 1 / batch_size # Configure input @@ -144,13 +154,15 @@ def log(x): optimizer_D.step() # Calculate loss of generator and update - g_loss = g_target * (torch.sum(d_real) + torch.sum(d_fake)) + log(Z) + g_loss = g_target * (torch.sum(d_real) + torch.sum(d_fake)) + log(Z) g_loss.backward() optimizer_G.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()) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) diff --git a/implementations/wgan/wgan.py b/implementations/wgan/wgan.py index 51a5d0b4..6c9008a7 100644 --- a/implementations/wgan/wgan.py +++ b/implementations/wgan/wgan.py @@ -15,19 +15,19 @@ 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=64, help='size of the batches') -parser.add_argument('--lr', type=float, default=0.00005, help='learning rate') -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=28, help='size of each image dimension') -parser.add_argument('--channels', type=int, default=1, help='number of image channels') -parser.add_argument('--n_critic', type=int, default=5, help='number of training steps for discriminator per iter') -parser.add_argument('--clip_value', type=float, default=0.01, help='lower and upper clip value for disc. weights') -parser.add_argument('--sample_interval', type=int, default=400, help='interval betwen image samples') +parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training") +parser.add_argument("--batch_size", type=int, default=64, help="size of the batches") +parser.add_argument("--lr", type=float, default=0.00005, help="learning rate") +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=28, help="size of each image dimension") +parser.add_argument("--channels", type=int, default=1, help="number of image channels") +parser.add_argument("--n_critic", type=int, default=5, help="number of training steps for discriminator per iter") +parser.add_argument("--clip_value", type=float, default=0.01, help="lower and upper clip value for disc. weights") +parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples") opt = parser.parse_args() print(opt) @@ -35,12 +35,13 @@ cuda = True if torch.cuda.is_available() else False + class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() def block(in_feat, out_feat, normalize=True): - layers = [ nn.Linear(in_feat, out_feat)] + layers = [nn.Linear(in_feat, out_feat)] if normalize: layers.append(nn.BatchNorm1d(out_feat, 0.8)) layers.append(nn.LeakyReLU(0.2, inplace=True)) @@ -60,6 +61,7 @@ def forward(self, z): img = img.view(img.shape[0], *img_shape) return img + class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() @@ -69,7 +71,7 @@ def __init__(self): nn.LeakyReLU(0.2, inplace=True), nn.Linear(512, 256), nn.LeakyReLU(0.2, inplace=True), - nn.Linear(256, 1) + nn.Linear(256, 1), ) def forward(self, img): @@ -77,6 +79,7 @@ def forward(self, img): validity = self.model(img_flat) return validity + # Initialize generator and discriminator generator = Generator() discriminator = Discriminator() @@ -86,14 +89,17 @@ def forward(self, img): discriminator.cuda() # Configure data loader -os.makedirs('../../data/mnist', exist_ok=True) +os.makedirs("../../data/mnist", exist_ok=True) dataloader = torch.utils.data.DataLoader( - datasets.MNIST('../../data/mnist', train=True, download=True, - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ])), - batch_size=opt.batch_size, shuffle=True) + datasets.MNIST( + "../../data/mnist", + train=True, + download=True, + transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]), + ), + batch_size=opt.batch_size, + shuffle=True, +) # Optimizers optimizer_G = torch.optim.RMSprop(generator.parameters(), lr=opt.lr) @@ -151,10 +157,11 @@ def forward(self, img): loss_G.backward() optimizer_G.step() - print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" % (epoch, opt.n_epochs, - batches_done % len(dataloader), len(dataloader), - loss_D.item(), loss_G.item())) + print( + "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" + % (epoch, opt.n_epochs, batches_done % len(dataloader), len(dataloader), loss_D.item(), loss_G.item()) + ) 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[:25], "images/%d.png" % batches_done, nrow=5, normalize=True) batches_done += 1 diff --git a/implementations/wgan_div/wgan_div.py b/implementations/wgan_div/wgan_div.py index 1b51f188..1b0a8096 100644 --- a/implementations/wgan_div/wgan_div.py +++ b/implementations/wgan_div/wgan_div.py @@ -102,7 +102,9 @@ def forward(self, img): "../../data/mnist", train=True, download=True, - transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]), + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), ), batch_size=opt.batch_size, shuffle=True, @@ -123,7 +125,7 @@ def forward(self, img): for i, (imgs, _) in enumerate(dataloader): # Configure input - real_imgs = Variable(imgs.type(Tensor),requires_grad=True) + real_imgs = Variable(imgs.type(Tensor), requires_grad=True) # --------------------- # Train Discriminator @@ -143,23 +145,17 @@ def forward(self, img): fake_validity = discriminator(fake_imgs) # Compute W-div gradient penalty - real_grad_out = Variable(Tensor(real_imgs.size(0), 1).fill_(1.0),requires_grad=False) - real_grad = autograd.grad(real_validity, - real_imgs, - real_grad_out, - create_graph=True, - retain_graph=True, - only_inputs=True)[0] - real_grad_norm = real_grad.view(real_grad.size(0),-1).pow(2).sum(1)**(p/2) - - fake_grad_out = Variable(Tensor(fake_imgs.size(0), 1).fill_(1.0),requires_grad=False) - fake_grad = autograd.grad(fake_validity, - fake_imgs, - fake_grad_out, - create_graph=True, - retain_graph=True, - only_inputs=True)[0] - fake_grad_norm = fake_grad.view(fake_grad.size(0),-1).pow(2).sum(1)**(p/2) + real_grad_out = Variable(Tensor(real_imgs.size(0), 1).fill_(1.0), requires_grad=False) + real_grad = autograd.grad( + real_validity, real_imgs, real_grad_out, create_graph=True, retain_graph=True, only_inputs=True + )[0] + real_grad_norm = real_grad.view(real_grad.size(0), -1).pow(2).sum(1) ** (p / 2) + + fake_grad_out = Variable(Tensor(fake_imgs.size(0), 1).fill_(1.0), requires_grad=False) + fake_grad = autograd.grad( + fake_validity, fake_imgs, fake_grad_out, create_graph=True, retain_graph=True, only_inputs=True + )[0] + fake_grad_norm = fake_grad.view(fake_grad.size(0), -1).pow(2).sum(1) ** (p / 2) div_gp = torch.mean(real_grad_norm + fake_grad_norm) * k / 2 diff --git a/implementations/wgan_gp/wgan_gp.py b/implementations/wgan_gp/wgan_gp.py index 2c5ace40..17d5afbf 100644 --- a/implementations/wgan_gp/wgan_gp.py +++ b/implementations/wgan_gp/wgan_gp.py @@ -101,7 +101,9 @@ def forward(self, img): "../../data/mnist", train=True, download=True, - transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]), + transform=transforms.Compose( + [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])] + ), ), batch_size=opt.batch_size, shuffle=True,