class CelebA(object): """Implement DCGAN for CelebA dataset""" def __init__(self, train_params, ckpt_params, gan_params): # Training parameters self.root_dir = train_params['root_dir'] self.gen_dir = train_params['gen_dir'] self.batch_size = train_params['batch_size'] self.train_len = train_params['train_len'] self.learning_rate = train_params['learning_rate'] self.momentum = train_params['momentum'] self.optim = train_params['optim'] self.use_cuda = train_params['use_cuda'] # Checkpoint parameters (when, where) self.batch_report_interval = ckpt_params['batch_report_interval'] self.ckpt_path = ckpt_params['ckpt_path'] self.save_stats_interval = ckpt_params['save_stats_interval'] # Create directories if they don't exist if not os.path.isdir(self.ckpt_path): os.mkdir(self.ckpt_path) if not os.path.isdir(self.gen_dir): os.mkdir(self.gen_dir) # GAN parameters self.gan_type = gan_params['gan_type'] self.latent_dim = gan_params['latent_dim'] self.n_critic = gan_params['n_critic'] # Make sure report interval divides total num of batches self.num_batches = self.train_len // self.batch_size # Get ready to ruuummmmmmble self.compile() def compile(self): """Compile model (loss function, optimizers, etc.)""" # Create new GAN self.gan = DCGAN(self.gan_type, self.latent_dim, self.batch_size, self.use_cuda) # Set optimizers for generator and discriminator if self.optim == 'adam': self.G_optimizer = optim.Adam(self.gan.G.parameters(), lr=self.learning_rate, betas=self.momentum) self.D_optimizer = optim.Adam(self.gan.D.parameters(), lr=self.learning_rate, betas=self.momentum) elif self.optim == 'rmsprop': self.G_optimizer = optim.RMSprop(self.gan.G.parameters(), lr=self.learning_rate) self.D_optimizer = optim.RMSprop(self.gan.D.parameters(), lr=self.learning_rate) else: raise NotImplementedError # CUDA support if torch.cuda.is_available() and self.use_cuda: self.gan = self.gan.cuda() # Create fixed latent variables for inference while training self.latent_vars = [] for i in range(100): self.latent_vars.append(self.gan.create_latent_var(1)) def save_stats(self, stats): """Save model statistics""" fname_pkl = '{}/{}-stats.pkl'.format(self.ckpt_path, self.gan_type) print('Saving model statistics to: {}'.format(fname_pkl)) with open(fname_pkl, 'wb') as fp: pickle.dump(stats, fp) def eval(self, n, epoch=None, while_training=False): """Sample examples from generator's distribution""" # Evaluation mode self.gan.G.eval() # Montage size (square) m = int(np.sqrt(n)) # Predict images to see progress for i in range(n): # Reuse fixed latent variables to keep random process intact if while_training: img = self.gan.generate_img(self.latent_vars[i]) else: img = self.gan.generate_img() img = utils.unnormalize(img.squeeze()) fname_in = '{}/test{:d}.png'.format(self.ckpt_path, i) torchvision.utils.save_image(img, fname_in) stack = 'montage {}/test* -tile {}x{} -geometry 64x64+1+1 \ {}/epoch'.format(self.ckpt_path, m, m, self.ckpt_path) stack = stack + str( epoch + 1) + '.png' if epoch is not None else stack + '.png' sp.call(stack.split()) for f in glob.glob('{}/test*'.format(self.ckpt_path)): os.remove(f) def train(self, nb_epochs, data_loader): """Train model on data""" # Initialize tracked quantities and prepare everything G_all_losses, D_all_losses, times = [], [], utils.AvgMeter() utils.format_hdr(self.gan, self.root_dir, self.train_len) start = datetime.datetime.now() g_iter, d_iter = 0, 0 # Train for epoch in range(nb_epochs): print('EPOCH {:d} / {:d}'.format(epoch + 1, nb_epochs)) G_losses, D_losses = utils.AvgMeter(), utils.AvgMeter() start_epoch = datetime.datetime.now() avg_time_per_batch = utils.AvgMeter() # Mini-batch SGD for batch_idx, (x, _) in enumerate(data_loader): # Critic update ratio if self.gan_type == 'wgan': n_critic = 20 if g_iter < 50 or ( g_iter + 1) % 500 == 0 else self.n_critic else: n_critic = self.n_critic # Training mode self.gan.G.train() # Discard last examples to simplify code if x.size(0) != self.batch_size: break batch_start = datetime.datetime.now() # Print progress bar utils.progress_bar(batch_idx, self.batch_report_interval, G_losses.avg, D_losses.avg) x = Variable(x) if torch.cuda.is_available() and self.use_cuda: x = x.cuda() # Update discriminator D_loss, fake_imgs = self.gan.train_D(x, self.D_optimizer, self.batch_size) D_losses.update(D_loss, self.batch_size) d_iter += 1 # Update generator if batch_idx % n_critic == 0: G_loss = self.gan.train_G(self.G_optimizer, self.batch_size) G_losses.update(G_loss, self.batch_size) g_iter += 1 batch_end = datetime.datetime.now() batch_time = int( (batch_end - batch_start).total_seconds() * 1000) avg_time_per_batch.update(batch_time) # Report model statistics if (batch_idx % self.batch_report_interval == 0 and batch_idx) or \ self.batch_report_interval == self.num_batches: G_all_losses.append(G_losses.avg) D_all_losses.append(D_losses.avg) utils.show_learning_stats(batch_idx, self.num_batches, G_losses.avg, D_losses.avg, avg_time_per_batch.avg) [ k.reset() for k in [G_losses, D_losses, avg_time_per_batch] ] self.eval(100, epoch=epoch, while_training=True) # print('Critic iter: {}'.format(g_iter)) # Save stats if batch_idx % self.save_stats_interval == 0 and batch_idx: stats = dict(G_loss=G_all_losses, D_loss=D_all_losses) self.save_stats(stats) # Save model utils.clear_line() print('Elapsed time for epoch: {}'.format( utils.time_elapsed_since(start_epoch))) self.gan.save_model(self.ckpt_path, epoch) self.eval(100, epoch=epoch, while_training=True) # Print elapsed time elapsed = utils.time_elapsed_since(start) print('Training done! Total elapsed time: {}\n'.format(elapsed)) return G_loss, D_loss
class CelebA(object): """Implement DCGAN for CelebA dataset""" def __init__(self, train_params, ckpt_params, gan_params): # Training parameters self.root_dir = train_params['root_dir'] self.batch_size = train_params['batch_size'] self.train_len = train_params['train_len'] self.learning_rate = train_params['learning_rate'] self.momentum = train_params['momentum'] self.optim = train_params['optim'] self.use_cuda = train_params['use_cuda'] # Checkpoint parameters (when, where) self.batch_report_interval = ckpt_params['batch_report_interval'] self.ckpt_path = ckpt_params['ckpt_path'] self.save_stats_interval = ckpt_params['save_stats_interval'] # Create directories if they don't exist if not os.path.isdir(self.ckpt_path): print(self.ckpt_path) os.mkdir(self.ckpt_path) # GAN parameters self.gan_type = gan_params['gan_type'] self.latent_dim = gan_params['latent_dim'] self.n_critic = gan_params['n_critic'] # Make sure report interval divides total num of batches self.num_batches = self.train_len // self.batch_size self.compile() #frequency weight self.freq_weight = 0 def compile(self): """Compile model (loss function, optimizers, etc.)""" # Create new GAN self.gan = DCGAN(self.gan_type, self.latent_dim, self.batch_size, self.use_cuda) # Set optimizers for generator and discriminator if self.optim == 'adam': self.G_optimizer = optim.Adam(self.gan.G.parameters(), lr=self.learning_rate, betas=self.momentum) self.D_optimizer = optim.Adam(self.gan.D.parameters(), lr=self.learning_rate, betas=self.momentum) elif self.optim == 'rmsprop': self.G_optimizer = optim.RMSprop(self.gan.G.parameters(), lr=self.learning_rate) self.D_optimizer = optim.RMSprop(self.gan.D.parameters(), lr=self.learning_rate) else: raise NotImplementedError # CUDA support if torch.cuda.is_available() and self.use_cuda: self.gan = self.gan.cuda() def save_stats(self, stats): """Save model statistics""" fname_pkl = '{}/{}-stats.pkl'.format(self.ckpt_path, self.gan_type) print('Saving model statistics to: {}'.format(fname_pkl)) with open(fname_pkl, 'wb') as fp: pickle.dump(stats, fp) def test(self, epoch): fname_gen_pt = '{}/{}-gen-epoch-{}.pt'.format(self.ckpt_path, self.gan_type, epoch + 1) self.gan.load_model(fname_gen_pt) directory = self.ckpt_path + "/testing/" + str(epoch + 1) if not os.path.exists(directory): os.makedirs(directory) # Evaluation mode self.gan.G.eval() n = 10000 # Predict images to see progress for i in range(n): img = self.gan.generate_img() img = utils.unnormalize(img.squeeze()) fname_in = '{}/{:d}_test.png'.format(directory, i) torchvision.utils.save_image(img, fname_in) def train(self, nb_epochs, data_loader): """Train model on data""" # Initialize tracked quantities and prepare everything G_all_losses, D_all_losses, times = [], [], utils.AvgMeter() utils.format_hdr(self.gan, self.root_dir, self.train_len) start = datetime.datetime.now() g_iter, d_iter = 0, 0 # Train for epoch in range(nb_epochs): print('EPOCH {:d} / {:d}'.format(epoch + 1, nb_epochs)) G_losses, D_losses = utils.AvgMeter(), utils.AvgMeter() start_epoch = datetime.datetime.now() avg_time_per_batch = utils.AvgMeter() # Mini-batch SGD for batch_idx, (x, _) in enumerate(data_loader): # Critic update ratio if self.gan_type == 'wgan': n_critic = 20 if g_iter < 50 or ( g_iter + 1) % 500 == 0 else self.n_critic else: n_critic = self.n_critic # Training mode self.gan.G.train() # Discard last examples to simplify code if x.size(0) != self.batch_size: break batch_start = datetime.datetime.now() # Print progress bar utils.progress_bar(batch_idx, self.batch_report_interval, G_losses.avg, D_losses.avg) x = Variable(x) if torch.cuda.is_available() and self.use_cuda: x = x.cuda() self.freq_weight = (epoch + 1) / nb_epochs # Update discriminator D_loss, fake_imgs = self.gan.train_D(x, self.freq_weight, self.D_optimizer, self.batch_size) D_losses.update(D_loss, self.batch_size) d_iter += 1 # Update generator if batch_idx % n_critic == 0: G_loss = self.gan.train_G(self.freq_weight, self.G_optimizer, self.batch_size) G_losses.update(G_loss, self.batch_size) g_iter += 1 batch_end = datetime.datetime.now() batch_time = int( (batch_end - batch_start).total_seconds() * 1000) avg_time_per_batch.update(batch_time) # Report model statistics if (batch_idx % self.batch_report_interval == 0 and batch_idx) or \ self.batch_report_interval == self.num_batches: G_all_losses.append(G_losses.avg) D_all_losses.append(D_losses.avg) utils.show_learning_stats(batch_idx, self.num_batches, G_losses.avg, D_losses.avg, avg_time_per_batch.avg) [ k.reset() for k in [G_losses, D_losses, avg_time_per_batch] ] # Save stats if batch_idx % self.save_stats_interval == 0 and batch_idx: stats = dict(G_loss=G_all_losses, D_loss=D_all_losses) self.save_stats(stats) # Save model utils.clear_line() print('Elapsed time for epoch: {}'.format( utils.time_elapsed_since(start_epoch))) self.gan.save_model(self.ckpt_path, epoch, False) # Generating model.test(epoch) # Print elapsed time elapsed = utils.time_elapsed_since(start) print('Training done! Total elapsed time: {}\n'.format(elapsed)) return G_loss, D_loss