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()) self.cuda = False self.cuda_index = 0 # 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)) if self.cuda: real_labels = Variable(torch.ones(self.batch_size)).cuda( self.cuda_index) fake_labels = Variable(torch.zeros(self.batch_size)).cuda( self.cuda_index) images, z = Variable(images.cuda( self.cuda_index)), Variable(z.cuda(self.cuda_index)) else: real_labels = Variable(torch.ones(self.batch_size)) fake_labels = Variable(torch.zeros(self.batch_size)) 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.flatten(), 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.flatten(), 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.flatten(), 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, g_loss.data)) if self.cuda: z = Variable( torch.randn(self.batch_size, 100).cuda(self.cuda_index)) else: z = Variable(torch.randn(self.batch_size, 100)) # ============ TensorBoard logging ============# # (1) Log the scalar values info = {'d_loss': d_loss.data, 'g_loss': g_loss.data} 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 if self.cuda: z = Variable( torch.randn(self.batch_size, 100).cuda(self.cuda_index)) else: z = Variable(torch.randn(self.batch_size, 100)) 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) if self.cuda: z = Variable( torch.randn(self.batch_size, 100).cuda(self.cuda_index)) else: z = Variable(torch.randn(self.batch_size, 100)) 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 DCGAN_MODEL(object): def __init__(self, args): print("DCGAN model initalization.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.C = args.channels # binary cross entropy loss and optimizer self.loss = nn.BCELoss() 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.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.loss = nn.BCELoss().cuda(self.cuda_index) 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) # 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=32, # resize=True, splits=10) print('Epoch-{}'.format(epoch + 1)) self.save_model(generator_iter) if not os.path.exists('dcgan_conv_training_result_images/'): os.makedirs('dcgan_conv_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, 'dcgan_conv_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(generator_iter) 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_conv_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, g_iter): torch.save(self.G.state_dict(), './checkpoint' + "/{}_DCGAN_CONV_G.pth".format(g_iter + 1)) torch.save(self.D.state_dict(), './checkpoint' + "/{}_DCGAN_CONV_D.pth".format(g_iter + 1)) print('Models save to ./checkpoint {}_G.pth & ./checkpoint {}_D.pth ') 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 DCGAN_MODEL(object): def __init__(self, args): print("DCGAN model initalization.") self.G = Generator(args.channels) self.D = Discriminator(args.channels) self.C = args.channels self.mode = args.mode self.name = ('res/_mode_' + str(args.mode) + '_beta_g_' + str(args.beta_g) + '_beta_g_' + str(args.beta_g) + '_beta_d_' + str(args.beta_d) + '_lr_g_' + str(args.lr_g) + '_lr_d_' + str(args.lr_d) + '_alpha_d_vjp_' + str(args.alpha_d_vjp) + '_alpha_g_vjp_' + str(args.alpha_g_vjp) + '_alpha_d_grad_' + str(args.alpha_d_grad) + '_alpha_g_grad_' + str(args.alpha_g_grad)) print(self.name) if not os.path.exists(self.name): os.makedirs(self.name) # binary cross entropy loss and optimizer self.loss = nn.BCEWithLogitsLoss() 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] if self.mode == 'adam': 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)) elif self.mode == 'adam_vjp': self.d_optimizer = optim.VJP_Adam(self.D.parameters(), lr=args.lr_d, betas=(args.beta_d, 0.999), alpha_vjp=args.alpha_d_vjp, alpha_grad=args.alpha_d_grad) self.g_optimizer = optim.VJP_Adam(self.G.parameters(), lr=args.lr_g, betas=(args.beta_g, 0.999), alpha_vjp=args.alpha_g_vjp, alpha_grad=args.alpha_g_grad) 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.loss = nn.BCEWithLogitsLoss().cuda(self.cuda_index) 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") dis_params_flatten = parameters_to_vector(self.D.parameters()) gen_params_flatten = parameters_to_vector(self.G.parameters()) # just to fill the empty grad buffers 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) fake_labels = torch.zeros(self.batch_size) fake_labels = Variable(fake_labels).cuda(self.cuda_index) d_loss_fake = self.loss(outputs.squeeze(), fake_labels) (0.0 * d_loss_fake).backward(create_graph=True) d_loss_fake = 0.0 best_inception_score = 0.0 d_loss_list = [] g_loss_list = [] 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) # Train discriminator # Compute BCE_Loss using real images outputs = self.D(images) d_loss_real = self.loss(outputs.squeeze(), 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.squeeze(), fake_labels) fake_score = outputs # Optimize discriminator d_loss = d_loss_real + d_loss_fake if self.mode == 'adam': self.D.zero_grad() d_loss.backward() self.d_optimizer.step() elif self.mode == 'adam_vjp': gradsD = torch.autograd.grad(outputs=d_loss, inputs=(self.D.parameters()), create_graph=True) for p, g in zip(self.D.parameters(), gradsD): p.grad = g gen_params_flatten_prev = gen_params_flatten + 0.0 gen_params_flatten = parameters_to_vector( self.G.parameters()) + 0.0 grad_gen_params_flatten = optim.parameters_grad_to_vector( self.G.parameters()) delta_gen_params_flatten = gen_params_flatten - gen_params_flatten_prev vjp_dis = torch.autograd.grad( grad_gen_params_flatten, self.D.parameters(), grad_outputs=delta_gen_params_flatten) self.d_optimizer.step(vjps=vjp_dis) # 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) # non-zero_sum g_loss = self.loss(outputs.squeeze(), real_labels) # zer_sum: # g_loss = - self.loss(outputs.squeeze(), fake_labels) # Optimize generator if self.mode == 'adam': self.D.zero_grad() self.G.zero_grad() g_loss.backward() self.g_optimizer.step() elif self.mode == 'adam_vjp': gradsG = torch.autograd.grad(outputs=g_loss, inputs=(self.G.parameters()), create_graph=True) for p, g in zip(self.G.parameters(), gradsG): p.grad = g dis_params_flatten_prev = dis_params_flatten + 0.0 dis_params_flatten = parameters_to_vector( self.D.parameters()) + 0.0 grad_dis_params_flatten = optim.parameters_grad_to_vector( self.D.parameters()) delta_dis_params_flatten = dis_params_flatten - dis_params_flatten_prev vjp_gen = torch.autograd.grad( grad_dis_params_flatten, self.G.parameters(), grad_outputs=delta_dis_params_flatten) self.g_optimizer.step(vjps=vjp_gen) 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=32, resize=True, splits=10) print('Epoch-{}'.format(epoch + 1)) print(inception_score) if inception_score >= best_inception_score: best_inception_score = inception_score self.save_model() # 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.name + '/iter_{}_inception_{}_.png'.format( str(generator_iter).zfill(3), str(inception_score))) 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: d_loss_list += [d_loss.item()] g_loss_list += [g_loss.item()] print("Epoch: [%2d] [%4d/%4d] D_loss: %.8f, G_loss: %.8f" % ((epoch + 1), (i + 1), train_loader.dataset.__len__() // self.batch_size, d_loss.item(), g_loss.item())) 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.item(), 'g_loss': g_loss.item()} 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))) # Save the trained parameters self.save_final_model() np.save(self.name + '/d_loss', np.array(d_loss_list)) np.save(self.name + '/g_loss', np.array(g_loss_list)) self.evaluate(train_loader, self.name + '/discriminator.pkl', self.name + '/generator.pkl') def evaluate(self, test_loader, D_model_path, G_model_path): self.load_model(D_model_path, G_model_path) self.G.eval() all_fake = [] for i in range(10): z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index) fake = self.G(z) all_fake += [fake] all_fake = torch.cat(all_fake, 0) inception_score = calc_inception_score( (all_fake.cpu().data.numpy() + 1.0) * 128) print(inception_score) 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.name + '/best_inception_score' + str(inception_score) + '.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(), self.name + '/generator.pkl') torch.save(self.D.state_dict(), self.name + '/discriminator.pkl') print('Models save to generator.pkl & discriminator.pkl ') def save_final_model(self): torch.save(self.G.state_dict(), self.name + '/final_generator.pkl') torch.save(self.D.state_dict(), self.name + '/final_discriminator.pkl') print('Final 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)))