def __init__(self, args): print("WGAN_CP init model.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.C = args.channels # check if cuda is available self.check_cuda(args.cuda) # WGAN values from paper self.learning_rate = 0.00005 self.batch_size = 64 self.weight_cliping_limit = 0.01 # WGAN with gradient clipping uses RMSprop instead of ADAM self.d_optimizer = torch.optim.RMSprop(self.D.parameters(), lr=self.learning_rate) self.g_optimizer = torch.optim.RMSprop(self.G.parameters(), lr=self.learning_rate) # Set the logger self.logger = Logger('./logs') self.logger.writer.flush() self.number_of_images = 10 self.generator_iters = args.generator_iters self.critic_iter = 5
def __init__(self, args): print("VLGAN model initalization.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.E = Encoder(args.channels) self.C = args.channels self.recon_weight = 5 # binary cross entropy loss and optimizer self.loss = nn.BCELoss() self.v_loss = Variational_loss self.cuda = "False" self.cuda_index = 0 # check if cuda is available self.check_cuda(args.cuda) # Using lower learning rate than suggested by (ADAM authors) lr=0.0002 and Beta_1 = 0.5 instead od 0.9 works better [Radford2015] self.d_optimizer = torch.optim.Adam(self.D.parameters(), lr=0.0002, betas=(0.5, 0.999)) self.g_optimizer = torch.optim.Adam(self.G.parameters(), lr=0.0002, betas=(0.5, 0.999)) self.e_optimizer = torch.optim.Adam(self.E.parameters(), lr=0.0002, betas=(0.5, 0.999)) self.epochs = args.epochs self.batch_size = args.batch_size # Set the logger self.logger = Logger('./logs') self.number_of_images = 10
def __init__(self, args): # Generator architecture self.G = nn.Sequential(nn.Linear(100, 256), nn.LeakyReLU(0.2), nn.Linear(256, 512), nn.LeakyReLU(0.2), nn.Linear(512, 1024), nn.LeakyReLU(0.2), nn.Tanh()) # Discriminator architecture self.D = nn.Sequential(nn.Linear(1024, 512), nn.LeakyReLU(0.2), nn.Linear(512, 256), nn.LeakyReLU(0.2), nn.Linear(256, 1), nn.Sigmoid()) # Check if cuda is available self.check_cuda(args.cuda) # Binary cross entropy loss and optimizer self.loss = nn.BCELoss() self.d_optimizer = torch.optim.Adam(self.D.parameters(), lr=0.0002, weight_decay=0.00001) self.g_optimizer = torch.optim.Adam(self.G.parameters(), lr=0.0002, weight_decay=0.00001) # Set the logger self.logger = Logger('./logs') self.number_of_images = 10 self.epochs = args.epochs self.batch_size = args.batch_size
def __init__(self, args): print("WGAN_GradientPenalty init model.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.C = args.channels # Check if cuda is available self.check_cuda(args.cuda) # WGAN values from paper self.learning_rate = 1e-4 self.b1 = 0.5 self.b2 = 0.999 self.batch_size = args.batch_size # WGAN_gradient penalty uses ADAM self.d_optimizer = optim.Adam(self.D.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2)) self.g_optimizer = optim.Adam(self.G.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2)) # Set the logger self.logger = Logger('./logs') self.logger.writer.flush() self.number_of_images = 10 self.generator_iters = args.generator_iters self.critic_iter = 5 self.lambda_term = 10
class WGAN_CP(object): def __init__(self, args): print("WGAN_CP init model.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.C = args.channels # check if cuda is available self.check_cuda(args.cuda) # WGAN values from paper self.learning_rate = 0.00005 self.batch_size = 64 self.weight_cliping_limit = 0.01 # WGAN with gradient clipping uses RMSprop instead of ADAM self.d_optimizer = torch.optim.RMSprop(self.D.parameters(), lr=self.learning_rate) self.g_optimizer = torch.optim.RMSprop(self.G.parameters(), lr=self.learning_rate) # Set the logger self.logger = Logger('./logs') self.logger.writer.flush() self.number_of_images = 10 self.generator_iters = args.generator_iters self.critic_iter = 5 def check_cuda(self, cuda_flag=False): if cuda_flag: self.cuda_index = 0 self.cuda = True self.D.cuda() self.G.cuda() print("Cuda enabled flag: {}".format(self.cuda)) def train(self, train_loader): self.t_begin = t.time() #self.file = open("inception_score_graph.txt", "w") # Now batches are callable self.data.next() self.data = self.get_infinite_batches(train_loader) one = torch.FloatTensor([1]) mone = one * -1 if self.cuda: one = one.cuda() mone = mone.cuda() for g_iter in range(self.generator_iters): # Requires grad, Generator requires_grad = False for p in self.D.parameters(): p.requires_grad = True # Train Dicriminator forward-loss-backward-update self.critic_iter times while 1 Generator forward-loss-backward-update for d_iter in range(self.critic_iter): self.D.zero_grad() # Clamp parameters to a range [-c, c], c=self.weight_cliping_limit for p in self.D.parameters(): p.data.clamp_(-self.weight_cliping_limit, self.weight_cliping_limit) images = self.data.__next__() # Check for batch to have full batch_size if (images.size()[0] != self.batch_size): continue z = torch.rand((self.batch_size, 100, 1, 1)) if self.cuda: images, z = Variable(images.cuda()), Variable(z.cuda()) else: images, z = Variable(images), Variable(z) # Train discriminator # WGAN - Training discriminator more iterations than generator # Train with real images d_loss_real = self.D(images) d_loss_real = d_loss_real.mean(0).view(1) d_loss_real.backward(one) # Train with fake images if self.cuda: z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda() else: z = Variable(torch.randn(self.batch_size, 100, 1, 1)) fake_images = self.G(z) d_loss_fake = self.D(fake_images) d_loss_fake = d_loss_fake.mean(0).view(1) d_loss_fake.backward(mone) d_loss = d_loss_fake - d_loss_real Wasserstein_D = d_loss_real - d_loss_fake self.d_optimizer.step() # Generator update for p in self.D.parameters(): p.requires_grad = False # to avoid computation self.G.zero_grad() # Train generator # Compute loss with fake images z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda() fake_images = self.G(z) g_loss = self.D(fake_images) g_loss = g_loss.mean().mean(0).view(1) g_loss.backward(one) g_cost = -g_loss self.g_optimizer.step() # Saving model and sampling images every 1000th generator iterations if (g_iter) % 1000 == 0: self.save_model() # Workaround because graphic card memory can't store more than 830 examples in memory for generating image # Therefore doing loop and generating 800 examples and stacking into list of samples to get 8000 generated images # This way Inception score is more correct since there are different generated examples from every class of Inception model # sample_list = [] # for i in range(10): # z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index) # samples = self.G(z) # sample_list.append(samples.data.cpu().numpy()) # # # Flattening list of list into one list # new_sample_list = list(chain.from_iterable(sample_list)) # print("Calculating Inception Score over 8k generated images") # # Feeding list of numpy arrays # inception_score = get_inception_score(new_sample_list, cuda=True, batch_size=32, # resize=True, splits=10) if not os.path.exists('training_result_images/'): os.makedirs('training_result_images/') # Denormalize images and save them in grid 8x8 z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index) samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu()[:64] grid = utils.make_grid(samples) utils.save_image(grid, 'training_result_images/img_generatori_iter_{}.png'.format(str(g_iter).zfill(3))) # Testing time = t.time() - self.t_begin #print("Inception score: {}".format(inception_score)) print("Generator iter: {}".format(g_iter)) print("Time {}".format(time)) # Write to file inception_score, gen_iters, time #output = str(g_iter) + " " + str(time) + " " + str(inception_score[0]) + "\n" #self.file.write(output) # ============ TensorBoard logging ============# # (1) Log the scalar values info = { 'Wasserstein distance': Wasserstein_D.data[0], 'Loss D': d_loss.data[0], 'Loss G': g_cost.data[0], 'Loss D Real': d_loss_real.data[0], 'Loss D Fake': d_loss_fake.data[0] } for tag, value in info.items(): self.logger.scalar_summary(tag, value, g_iter + 1) # (3) Log the images info = { 'real_images': self.real_images(images, self.number_of_images), 'generated_images': self.generate_img(z, self.number_of_images) } for tag, images in info.items(): self.logger.image_summary(tag, images, g_iter + 1) self.t_end = t.time() print('Time of training-{}'.format((self.t_end - self.t_begin))) #self.file.close() # Save the trained parameters self.save_model() def evaluate(self, test_loader, D_model_path, G_model_path): self.load_model(D_model_path, G_model_path) z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda() samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu() grid = utils.make_grid(samples) print("Grid of 8x8 images saved to 'dgan_model_image.png'.") utils.save_image(grid, 'dgan_model_image.png') def real_images(self, images, number_of_images): if (self.C == 3): return self.to_np(images.view(-1, self.C, 32, 32)[:self.number_of_images]) else: return self.to_np(images.view(-1, 32, 32)[:self.number_of_images]) def generate_img(self, z, number_of_images): samples = self.G(z).data.cpu().numpy()[:number_of_images] generated_images = [] for sample in samples: if self.C == 3: generated_images.append(sample.reshape(self.C, 32, 32)) else: generated_images.append(sample.reshape(32, 32)) return generated_images def to_np(self, x): return x.data.cpu().numpy() def save_model(self): torch.save(self.G.state_dict(), './generator.pkl') torch.save(self.D.state_dict(), './discriminator.pkl') print('Models save to ./generator.pkl & ./discriminator.pkl ') def load_model(self, D_model_filename, G_model_filename): D_model_path = os.path.join(os.getcwd(), D_model_filename) G_model_path = os.path.join(os.getcwd(), G_model_filename) self.D.load_state_dict(torch.load(D_model_path)) self.G.load_state_dict(torch.load(G_model_path)) print('Generator model loaded from {}.'.format(G_model_path)) print('Discriminator model loaded from {}-'.format(D_model_path)) def get_infinite_batches(self, data_loader): while True: for i, (images, _) in enumerate(data_loader): yield images def generate_latent_walk(self, number): if not os.path.exists('interpolated_images/'): os.makedirs('interpolated_images/') number_int = 10 # interpolate between twe noise(z1, z2). z_intp = torch.FloatTensor(1, 100, 1, 1) z1 = torch.randn(1, 100, 1, 1) z2 = torch.randn(1, 100, 1, 1) if self.cuda: z_intp = z_intp.cuda() z1 = z1.cuda() z2 = z2.cuda() z_intp = Variable(z_intp) images = [] alpha = 1.0 / float(number_int + 1) print(alpha) for i in range(1, number_int + 1): z_intp.data = z1*alpha + z2*(1.0 - alpha) alpha += alpha fake_im = self.G(z_intp) fake_im = fake_im.mul(0.5).add(0.5) #denormalize images.append(fake_im.view(self.C,32,32).data.cpu()) grid = utils.make_grid(images, nrow=number_int ) utils.save_image(grid, 'interpolated_images/interpolated_{}.png'.format(str(number).zfill(3))) print("Saved interpolated images.")
class VLGAN_MODEL(object): def __init__(self, args): print("VLGAN model initalization.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.E = Encoder(args.channels) self.C = args.channels self.recon_weight = 5 # binary cross entropy loss and optimizer self.loss = nn.BCELoss() self.v_loss = Variational_loss self.cuda = "False" self.cuda_index = 0 # check if cuda is available self.check_cuda(args.cuda) # Using lower learning rate than suggested by (ADAM authors) lr=0.0002 and Beta_1 = 0.5 instead od 0.9 works better [Radford2015] self.d_optimizer = torch.optim.Adam(self.D.parameters(), lr=0.0002, betas=(0.5, 0.999)) self.g_optimizer = torch.optim.Adam(self.G.parameters(), lr=0.0002, betas=(0.5, 0.999)) self.e_optimizer = torch.optim.Adam(self.E.parameters(), lr=0.0002, betas=(0.5, 0.999)) self.epochs = args.epochs self.batch_size = args.batch_size # Set the logger self.logger = Logger('./logs') self.number_of_images = 10 # cuda support def check_cuda(self, cuda_flag=False): if cuda_flag: self.cuda = True self.D.cuda(self.cuda_index) self.G.cuda(self.cuda_index) self.E.cuda(self.cuda_index) self.loss = nn.BCELoss().cuda(self.cuda_index) L1_loss = nn.L1Loss().cuda(self.cuda_index) def Variational_loss(input, target, mu, logvar): alpha = 1 beta = 1 recon_loss = L1_loss(input, target) batch_size = logvar.data.shape[0] nz = logvar.data.shape[1] KLD_loss = (-0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())) / (nz * batch_size) return alpha * recon_loss, beta * KLD_loss self.v_loss = Variational_loss print("Cuda enabled flag: ") print(self.cuda) def train(self, train_loader): self.t_begin = t.time() generator_iter = 0 self.file = open("inception_score_graph.txt", "w") for epoch in range(self.epochs): self.epoch_start_time = t.time() for i, (images, _) in enumerate(train_loader): # Check if round number of batches if i == train_loader.dataset.__len__() // self.batch_size: break z = torch.rand((self.batch_size, 100, 1, 1)) real_labels = torch.ones(self.batch_size) fake_labels = torch.zeros(self.batch_size) if self.cuda: images, z = Variable(images).cuda( self.cuda_index), Variable(z).cuda(self.cuda_index) real_labels, fake_labels = Variable(real_labels).cuda( self.cuda_index), Variable(fake_labels).cuda( self.cuda_index) else: images, z = Variable(images), Variable(z) real_labels, fake_labels = Variable(real_labels), Variable( fake_labels) self.e_optimizer.zero_grad() self.g_optimizer.zero_grad() mu, logvar = self.E(images) std = torch.exp(0.5 * logvar) eps = Variable(torch.randn(std.size()), requires_grad=False).cuda(self.cuda_index) z = eps.mul(std).add_(mu) recon_image = self.G(z) err_recon, err_KLD = self.v_loss(recon_image, images.detach(), mu, logvar) err = self.recon_weight * (err_recon + err_KLD) err.backward() self.e_optimizer.step() self.g_optimizer.step() if self.cuda: images, z = Variable(images).cuda( self.cuda_index), Variable(z).cuda(self.cuda_index) real_labels, fake_labels = Variable(real_labels).cuda( self.cuda_index), Variable(fake_labels).cuda( self.cuda_index) else: images, z = Variable(images), Variable(z) real_labels, fake_labels = Variable(real_labels), Variable( fake_labels) # Train discriminator # Compute BCE_Loss using real images outputs = self.D(images) d_loss_real = self.loss(outputs, real_labels) real_score = outputs # Compute BCE Loss using fake images if self.cuda: z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index) else: z = Variable(torch.randn(self.batch_size, 100, 1, 1)) fake_images = self.G(z) outputs = self.D(fake_images) d_loss_fake = self.loss(outputs, fake_labels) fake_score = outputs # Optimize discriminator d_loss = d_loss_real + d_loss_fake self.D.zero_grad() d_loss.backward() self.d_optimizer.step() # Train generator # Compute loss with fake images if self.cuda: z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index) else: z = Variable(torch.randn(self.batch_size, 100, 1, 1)) fake_images = self.G(z) outputs = self.D(fake_images) g_loss = self.loss(outputs, real_labels) # Optimize generator self.D.zero_grad() self.G.zero_grad() g_loss.backward() self.g_optimizer.step() generator_iter += 1 if generator_iter % 1000 == 0: # Workaround because graphic card memory can't store more than 800+ examples in memory for generating image # Therefore doing loop and generating 800 examples and stacking into list of samples to get 8000 generated images # This way Inception score is more correct since there are different generated examples from every class of Inception model sample_list = [] for i in range(10): z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index) samples = self.G(z) sample_list.append(samples.data.cpu().numpy()) # Flattening list of lists into one list of numpy arrays new_sample_list = list(chain.from_iterable(sample_list)) print( "Calculating Inception Score over 8k generated images") # Feeding list of numpy arrays inception_score = get_inception_score(new_sample_list, cuda=True, batch_size=16, resize=True, splits=10) print('Epoch-{}'.format(epoch + 1)) self.save_model() if not os.path.exists('training_result_images/'): os.makedirs('training_result_images/') # Denormalize images and save them in grid 8x8 z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index) samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu()[:64] grid = utils.make_grid(samples) utils.save_image( grid, 'training_result_images/img_generatori_iter_{}.png'. format(str(generator_iter).zfill(3))) time = t.time() - self.t_begin print("Inception score: {}".format(inception_score)) print("Generator iter: {}".format(generator_iter)) print("Time {}".format(time)) # Write to file inception_score, gen_iters, time output = str(generator_iter) + " " + str(time) + " " + str( inception_score[0]) + "\n" self.file.write(output) if ((i + 1) % 100) == 0: print("Epoch: [%2d] [%4d/%4d] D_loss: %.8f, G_loss: %.8f" % ((epoch + 1), (i + 1), train_loader.dataset.__len__() // self.batch_size, d_loss.data[0], g_loss.data[0])) z = Variable( torch.randn(self.batch_size, 100, 1, 1).cuda(self.cuda_index)) # TensorBoard logging # Log the scalar values info = {'d_loss': d_loss.data[0], 'g_loss': g_loss.data[0]} for tag, value in info.items(): self.logger.scalar_summary(tag, value, generator_iter) # Log values and gradients of the parameters for tag, value in self.D.named_parameters(): tag = tag.replace('.', '/') self.logger.histo_summary(tag, self.to_np(value), generator_iter) self.logger.histo_summary(tag + '/grad', self.to_np(value.grad), generator_iter) # Log the images while training info = { 'real_images': self.real_images(images, self.number_of_images), 'generated_images': self.generate_img(z, self.number_of_images) } for tag, images in info.items(): self.logger.image_summary(tag, images, generator_iter) self.t_end = t.time() print('Time of training-{}'.format((self.t_end - self.t_begin))) #self.file.close() # Save the trained parameters self.save_model() def evaluate(self, test_loader, D_model_path, G_model_path): self.load_model(D_model_path, G_model_path) z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index) samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu() grid = utils.make_grid(samples) print("Grid of 8x8 images saved to 'dgan_model_image.png'.") utils.save_image(grid, 'dgan_model_image.png') def real_images(self, images, number_of_images): if (self.C == 3): return self.to_np( images.view(-1, self.C, 32, 32)[:self.number_of_images]) else: return self.to_np(images.view(-1, 32, 32)[:self.number_of_images]) def generate_img(self, z, number_of_images): samples = self.G(z).data.cpu().numpy()[:number_of_images] generated_images = [] for sample in samples: if self.C == 3: generated_images.append(sample.reshape(self.C, 32, 32)) else: generated_images.append(sample.reshape(32, 32)) return generated_images def to_np(self, x): return x.data.cpu().numpy() def save_model(self): torch.save(self.G.state_dict(), './generator.pkl') torch.save(self.D.state_dict(), './discriminator.pkl') print('Models save to ./generator.pkl & ./discriminator.pkl ') def load_model(self, D_model_filename, G_model_filename): D_model_path = os.path.join(os.getcwd(), D_model_filename) G_model_path = os.path.join(os.getcwd(), G_model_filename) self.D.load_state_dict(torch.load(D_model_path)) self.G.load_state_dict(torch.load(G_model_path)) print('Generator model loaded from {}.'.format(G_model_path)) print('Discriminator model loaded from {}-'.format(D_model_path)) def generate_latent_walk(self, number): if not os.path.exists('interpolated_images/'): os.makedirs('interpolated_images/') # Interpolate between twe noise(z1, z2) with number_int steps between number_int = 10 z_intp = torch.FloatTensor(1, 100, 1, 1) z1 = torch.randn(1, 100, 1, 1) z2 = torch.randn(1, 100, 1, 1) if self.cuda: z_intp = z_intp.cuda() z1 = z1.cuda() z2 = z2.cuda() z_intp = Variable(z_intp) images = [] alpha = 1.0 / float(number_int + 1) print(alpha) for i in range(1, number_int + 1): z_intp.data = z1 * alpha + z2 * (1.0 - alpha) alpha += alpha fake_im = self.G(z_intp) fake_im = fake_im.mul(0.5).add(0.5) #denormalize images.append(fake_im.view(self.C, 32, 32).data.cpu()) grid = utils.make_grid(images, nrow=number_int) utils.save_image( grid, 'interpolated_images/interpolated_{}.png'.format( str(number).zfill(3))) print( "Saved interpolated images to interpolated_images/interpolated_{}." .format(str(number).zfill(3)))
class GAN(object): def __init__(self, args): # Generator architecture self.G = nn.Sequential(nn.Linear(100, 256), nn.LeakyReLU(0.2), nn.Linear(256, 512), nn.LeakyReLU(0.2), nn.Linear(512, 1024), nn.LeakyReLU(0.2), nn.Tanh()) # Discriminator architecture self.D = nn.Sequential(nn.Linear(1024, 512), nn.LeakyReLU(0.2), nn.Linear(512, 256), nn.LeakyReLU(0.2), nn.Linear(256, 1), nn.Sigmoid()) # Check if cuda is available self.check_cuda(args.cuda) # Binary cross entropy loss and optimizer self.loss = nn.BCELoss() self.d_optimizer = torch.optim.Adam(self.D.parameters(), lr=0.0002, weight_decay=0.00001) self.g_optimizer = torch.optim.Adam(self.G.parameters(), lr=0.0002, weight_decay=0.00001) # Set the logger self.logger = Logger('./logs') self.number_of_images = 10 self.epochs = args.epochs self.batch_size = args.batch_size # Cuda support def check_cuda(self, cuda_flag=False): if cuda_flag: self.cuda_index = 0 self.cuda = True self.D.cuda(self.cuda_index) self.G.cuda(self.cuda_index) self.loss = nn.BCELoss().cuda(self.cuda_index) print("Cuda enabled flag: ") print(self.cuda) def train(self, train_loader): self.t_begin = time.time() generator_iter = 0 for epoch in range(self.epochs + 1): for i, (images, _) in enumerate(train_loader): # Check if round number of batches if i == train_loader.dataset.__len__() // self.batch_size: break # Flatten image 1,32x32 to 1024 images = images.view(self.batch_size, -1) z = torch.rand((self.batch_size, 100)) real_labels = Variable(torch.ones(self.batch_size)).cuda( self.cuda_index) fake_labels = Variable(torch.zeros(self.batch_size)).cuda( self.cuda_index) if self.cuda: images, z = Variable(images.cuda( self.cuda_index)), Variable(z.cuda(self.cuda_index)) else: images, z = Variable(images), Variable(z) # Train discriminator # compute BCE_Loss using real images where BCE_Loss(x, y): - y * log(D(x)) - (1-y) * log(1 - D(x)) # [Training discriminator = Maximizing discriminator being correct] outputs = self.D(images) d_loss_real = self.loss(outputs, real_labels) real_score = outputs # Compute BCELoss using fake images fake_images = self.G(z) outputs = self.D(fake_images) d_loss_fake = self.loss(outputs, fake_labels) fake_score = outputs # Optimizie discriminator d_loss = d_loss_real + d_loss_fake self.D.zero_grad() d_loss.backward() self.d_optimizer.step() # Train generator if self.cuda: z = Variable( torch.randn(self.batch_size, 100).cuda(self.cuda_index)) else: z = Variable(torch.randn(self.batch_size, 100)) fake_images = self.G(z) outputs = self.D(fake_images) # We train G to maximize log(D(G(z))[maximize likelihood of discriminator being wrong] instead of # minimizing log(1-D(G(z)))[minizing likelihood of discriminator being correct] # From paper [https://arxiv.org/pdf/1406.2661.pdf] g_loss = self.loss(outputs, real_labels) # Optimize generator self.D.zero_grad() self.G.zero_grad() g_loss.backward() self.g_optimizer.step() generator_iter += 1 if ((i + 1) % 100) == 0: print("Epoch: [%2d] [%4d/%4d] D_loss: %.8f, G_loss: %.8f" % ((epoch + 1), (i + 1), train_loader.dataset.__len__() // self.batch_size, d_loss.data[0], g_loss.data[0])) z = Variable( torch.randn(self.batch_size, 100).cuda(self.cuda_index)) # ============ TensorBoard logging ============# # (1) Log the scalar values info = {'d_loss': d_loss.data[0], 'g_loss': g_loss.data[0]} for tag, value in info.items(): self.logger.scalar_summary(tag, value, i + 1) # (2) Log values and gradients of the parameters (histogram) for tag, value in self.D.named_parameters(): tag = tag.replace('.', '/') self.logger.histo_summary(tag, self.to_np(value), i + 1) self.logger.histo_summary(tag + '/grad', self.to_np(value.grad), i + 1) # (3) Log the images info = { 'real_images': self.to_np( images.view(-1, 32, 32)[:self.number_of_images]), 'generated_images': self.generate_img(z, self.number_of_images) } for tag, images in info.items(): self.logger.image_summary(tag, images, i + 1) if generator_iter % 1000 == 0: print('Generator iter-{}'.format(generator_iter)) self.save_model() if not os.path.exists('training_result_images/'): os.makedirs('training_result_images/') # Denormalize images and save them in grid 8x8 z = Variable(torch.randn(self.batch_size, 100)).cuda(self.cuda_index) samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu() grid = utils.make_grid(samples) utils.save_image( grid, 'training_result_images/gan_image_iter_{}.png'.format( str(generator_iter).zfill(3))) self.t_end = time.time() print('Time of training-{}'.format((self.t_end - self.t_begin))) # Save the trained parameters self.save_model() def evaluate(self, test_loader, D_model_path, G_model_path): self.load_model(D_model_path, G_model_path) z = Variable(torch.randn(self.batch_size, 100)).cuda(self.cuda_index) samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu() grid = utils.make_grid(samples) print("Grid of 8x8 images saved to 'gan_model_image.png'.") utils.save_image(grid, 'gan_model_image.png') def generate_img(self, z, number_of_images): samples = self.G(z).data.cpu().numpy()[:number_of_images] generated_images = [] for sample in samples: generated_images.append(sample.reshape(32, 32)) return generated_images def to_np(self, x): return x.data.cpu().numpy() def save_model(self): torch.save(self.G.state_dict(), './generator.pkl') torch.save(self.D.state_dict(), './discriminator.pkl') print('Models save to ./generator.pkl & ./discriminator.pkl ') def load_model(self, D_model_filename, G_model_filename): D_model_path = os.path.join(os.getcwd(), D_model_filename) G_model_path = os.path.join(os.getcwd(), G_model_filename) self.D.load_state_dict(torch.load(D_model_path)) self.G.load_state_dict(torch.load(G_model_path)) print('Generator model loaded from {}.'.format(G_model_path)) print('Discriminator model loaded from {}-'.format(D_model_path))
class WGAN_GP(object): def __init__(self, args): print("WGAN_GradientPenalty init model.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.C = args.channels # Check if cuda is available self.check_cuda(args.cuda) # WGAN values from paper self.learning_rate = 1e-4 self.b1 = 0.5 self.b2 = 0.999 self.batch_size = args.batch_size # WGAN_gradient penalty uses ADAM self.d_optimizer = optim.Adam(self.D.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2)) self.g_optimizer = optim.Adam(self.G.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2)) # Set the logger self.logger = Logger('./logs') self.logger.writer.flush() self.number_of_images = 10 self.generator_iters = args.generator_iters self.critic_iter = 5 self.lambda_term = 10 def check_cuda(self, cuda_flag=False): if cuda_flag: self.cuda_index = 0 self.cuda = True self.D.cuda(self.cuda_index) self.G.cuda(self.cuda_index) print("Cuda enabled flag: {}".format(self.cuda)) def train(self, train_loader): self.t_begin = t.time() self.file = open(self.args.resultdir + "inception_score_graph.txt", "w") # Now batches are callable self.data.next() self.data = self.get_infinite_batches(train_loader) one = torch.tensor(1.0) mone = torch.tensor(-1.0) if self.cuda: one = one.cuda(self.cuda_index) mone = mone.cuda(self.cuda_index) for g_iter in range(self.generator_iters): print("[%d/%d]"%(g_iter, self.generator_iters)) # Requires grad, Generator requires_grad = False for p in self.D.parameters(): p.requires_grad = True d_loss_real = 0 d_loss_fake = 0 Wasserstein_D = 0 # Train Dicriminator forward-loss-backward-update self.critic_iter times while 1 Generator forward-loss-backward-update for d_iter in range(self.critic_iter): self.D.zero_grad() images = self.data.__next__() # Check for batch to have full batch_size if (images.size()[0] != self.batch_size): continue z = torch.rand((self.batch_size, 100, 1, 1)) if self.cuda: images, z = Variable(images.cuda(self.cuda_index)), Variable(z.cuda(self.cuda_index)) else: images, z = Variable(images), Variable(z) # Train discriminator # WGAN - Training discriminator more iterations than generator # Train with real images d_loss_real = self.D(images) d_loss_real = d_loss_real.mean() d_loss_real.backward(mone) # Train with fake images if self.cuda: z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index) else: z = Variable(torch.randn(self.batch_size, 100, 1, 1)) fake_images = self.G(z) d_loss_fake = self.D(fake_images) d_loss_fake = d_loss_fake.mean() d_loss_fake.backward(one) # Train with gradient penalty gradient_penalty = self.calculate_gradient_penalty(images.data, fake_images.data) gradient_penalty.backward(retain_graph=True) d_loss = d_loss_fake - d_loss_real + gradient_penalty Wasserstein_D = d_loss_real - d_loss_fake self.d_optimizer.step() # Generator update for p in self.D.parameters(): p.requires_grad = False # to avoid computation self.G.zero_grad() # train generator # compute loss with fake images z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index) fake_images = self.G(z) g_loss = self.D(fake_images) g_loss = g_loss.mean() g_loss.backward(mone) g_cost = -g_loss self.g_optimizer.step() with torch.no_grad(): # Saving model and sampling images every 1000th generator iterations if (g_iter) % 1000 == 0: self.save_model() # # Workaround because graphic card memory can't store more than 830 examples in memory for generating image # # Therefore doing loop and generating 800 examples and stacking into list of samples to get 8000 generated images # # This way Inception score is more correct since there are different generated examples from every class of Inception model sample_list = [] for i in range(10): #samples = self.data.__next__() z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index) samples = self.G(z) sample_list.append(samples.data.cpu().numpy()) # # # Flattening list of list into one list new_sample_list = list(chain.from_iterable(sample_list)) print("Calculating Inception Score over 8k generated images") # # Feeding list of numpy arrays inception_score = get_inception_score(new_sample_list, cuda=True, batch_size=32, resize=True, splits=10) if not os.path.exists(self.args.resultdir + 'training_result_images/'): os.makedirs(self.args.resultdir + 'training_result_images/') # Denormalize images and save them in grid 8x8 z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index) samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu()[:64] grid = utils.make_grid(samples) utils.save_image(grid, self.args.resultdir + 'training_result_images/img_generatori_iter_{}.png'.format(str(g_iter).zfill(3))) # Testing time = t.time() - self.t_begin #print("Real Inception score: {}".format(inception_score)) print("Generator iter: {}".format(g_iter)) print("Time {}".format(time)) # Write to file inception_score, gen_iters, time output = str(g_iter) + " " + str(time) + " " + str(inception_score[0]) + "\n" self.file.write(output) # ============ TensorBoard logging ============# # (1) Log the scalar valuesself.args.resultdir + info = { 'Wasserstein distance': Wasserstein_D.item(), 'Loss D': d_loss.item(), 'Loss G': g_cost.item(), 'Loss D Real': d_loss_real.item(), 'Loss D Fake': d_loss_fake.item() } for tag, value in info.items(): self.logger.scalar_summary(tag, value, g_iter + 1) # (3) Log the images info = { 'real_images': self.real_images(images, self.number_of_images), 'generated_images': self.generate_img(z, self.number_of_images) } for tag, images in info.items(): self.logger.image_summary(tag, images, g_iter + 1) self.t_end = t.time() print('Time of training-{}'.format((self.t_end - self.t_begin))) #self.file.close() # Save the trained parameters self.save_model() def evaluate(self, test_loader, D_model_path, G_model_path): self.load_model(D_model_path, G_model_path) z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index) samples = self.G(z) samples = samples.mul(0.5).add(0.5) samples = samples.data.cpu() grid = utils.make_grid(samples) print("Grid of 8x8 images saved to 'dgan_model_image.png'.") utils.save_image(grid, self.args.resultdir + 'dgan_model_image.png') def calculate_gradient_penalty(self, real_images, fake_images): eta = torch.FloatTensor(self.batch_size,1,1,1).uniform_(0,1) eta = eta.expand(self.batch_size, real_images.size(1), real_images.size(2), real_images.size(3)) if self.cuda: eta = eta.cuda(self.cuda_index) else: eta = eta interpolated = eta * real_images + ((1 - eta) * fake_images) if self.cuda: interpolated = interpolated.cuda(self.cuda_index) else: interpolated = interpolated # define it to calculate gradient interpolated = Variable(interpolated, requires_grad=True) # calculate probability of interpolated examples prob_interpolated = self.D(interpolated) # calculate gradients of probabilities with respect to examples gradients = autograd.grad(outputs=prob_interpolated, inputs=interpolated, grad_outputs=torch.ones( prob_interpolated.size()).cuda(self.cuda_index) if self.cuda else torch.ones( prob_interpolated.size()), create_graph=True, retain_graph=True)[0] grad_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * self.lambda_term return grad_penalty def real_images(self, images, number_of_images): if (self.C == 3): return self.to_np(images.view(-1, self.C, 32, 32)[:self.number_of_images]) else: return self.to_np(images.view(-1, 32, 32)[:self.number_of_images]) def generate_img(self, z, number_of_images): samples = self.G(z).data.cpu().numpy()[:number_of_images] generated_images = [] for sample in samples: if self.C == 3: generated_images.append(sample.reshape(self.C, 32, 32)) else: generated_images.append(sample.reshape(32, 32)) return generated_images def to_np(self, x): return x.data.cpu().numpy() def save_model(self): torch.save(self.G.state_dict(), './' + self.args.resultdir + 'generator.pkl') torch.save(self.D.state_dict(), './' + self.args.resultdir + 'discriminator.pkl') print('Models save to ./generator.pkl & ./discriminator.pkl ') def load_model(self, D_model_filename, G_model_filename): D_model_path = os.path.join(os.getcwd(), D_model_filename) G_model_path = os.path.join(os.getcwd(), G_model_filename) self.D.load_state_dict(torch.load(D_model_path)) self.G.load_state_dict(torch.load(G_model_path)) print('Generator model loaded from {}.'.format(G_model_path)) print('Discriminator model loaded from {}-'.format(D_model_path)) def get_infinite_batches(self, data_loader): while True: for i, (images, _) in enumerate(data_loader): yield images def generate_latent_walk(self, number): if not os.path.exists(self.args.resultdir + 'interpolated_images/'): os.makedirs(self.args.resultdir + 'interpolated_images/') number_int = 10 # interpolate between twe noise(z1, z2). z_intp = torch.FloatTensor(1, 100, 1, 1) z1 = torch.randn(1, 100, 1, 1) z2 = torch.randn(1, 100, 1, 1) if self.cuda: z_intp = z_intp.cuda() z1 = z1.cuda() z2 = z2.cuda() z_intp = Variable(z_intp) images = [] alpha = 1.0 / float(number_int + 1) print(alpha) for i in range(1, number_int + 1): z_intp.data = z1*alpha + z2*(1.0 - alpha) alpha += alpha fake_im = self.G(z_intp) fake_im = fake_im.mul(0.5).add(0.5) #denormalize images.append(fake_im.view(self.C,32,32).data.cpu()) grid = utils.make_grid(images, nrow=number_int ) utils.save_image(grid, self.args.resultdir + 'interpolated_images/interpolated_{}.png'.format(str(number).zfill(3))) print("Saved interpolated images.")