import os import random import matplotlib import matplotlib.pyplot as plt matplotlib.use('Agg') import torch from torch import nn, autograd from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter import torch.nn.functional as F from utils import common, train_utils from criteria import id_loss, moco_loss from configs import data_configs from datasets.images_dataset import ImagesDataset from criteria.lpips.lpips import LPIPS from models.psp import pSp from models.latent_codes_pool import LatentCodesPool from models.discriminator import LatentCodesDiscriminator from models.encoders.psp_encoders import ProgressiveStage from training.ranger import Ranger random.seed(0) torch.manual_seed(0) class Coach: def __init__(self, opts, prev_train_checkpoint=None): self.opts = opts self.global_step = 0 self.device = 'cuda:0' self.opts.device = self.device # Initialize network self.net = pSp(self.opts).to(self.device) # Initialize loss if self.opts.lpips_lambda > 0: self.lpips_loss = LPIPS(net_type=self.opts.lpips_type).to(self.device).eval() if self.opts.id_lambda > 0: if 'ffhq' in self.opts.dataset_type or 'celeb' in self.opts.dataset_type: self.id_loss = id_loss.IDLoss().to(self.device).eval() else: self.id_loss = moco_loss.MocoLoss(opts).to(self.device).eval() self.mse_loss = nn.MSELoss().to(self.device).eval() # Initialize optimizer self.optimizer = self.configure_optimizers() # Initialize discriminator if self.opts.w_discriminator_lambda > 0: self.discriminator = LatentCodesDiscriminator(512, 4).to(self.device) self.discriminator_optimizer = torch.optim.Adam(list(self.discriminator.parameters()), lr=opts.w_discriminator_lr) self.real_w_pool = LatentCodesPool(self.opts.w_pool_size) self.fake_w_pool = LatentCodesPool(self.opts.w_pool_size) # Initialize dataset self.train_dataset, self.test_dataset = self.configure_datasets() self.train_dataloader = DataLoader(self.train_dataset, batch_size=self.opts.batch_size, shuffle=True, num_workers=int(self.opts.workers), drop_last=True) self.test_dataloader = DataLoader(self.test_dataset, batch_size=self.opts.test_batch_size, shuffle=False, num_workers=int(self.opts.test_workers), drop_last=True) # Initialize logger log_dir = os.path.join(opts.exp_dir, 'logs') os.makedirs(log_dir, exist_ok=True) self.logger = SummaryWriter(log_dir=log_dir) # Initialize checkpoint dir self.checkpoint_dir = os.path.join(opts.exp_dir, 'checkpoints') os.makedirs(self.checkpoint_dir, exist_ok=True) self.best_val_loss = None if self.opts.save_interval is None: self.opts.save_interval = self.opts.max_steps if prev_train_checkpoint is not None: self.load_from_train_checkpoint(prev_train_checkpoint) prev_train_checkpoint = None def load_from_train_checkpoint(self, ckpt): print('Loading previous training data...') self.global_step = ckpt['global_step'] + 1 self.best_val_loss = ckpt['best_val_loss'] self.net.load_state_dict(ckpt['state_dict']) if self.opts.keep_optimizer: self.optimizer.load_state_dict(ckpt['optimizer']) if self.opts.w_discriminator_lambda > 0: self.discriminator.load_state_dict(ckpt['discriminator_state_dict']) self.discriminator_optimizer.load_state_dict(ckpt['discriminator_optimizer_state_dict']) if self.opts.progressive_steps: self.check_for_progressive_training_update(is_resume_from_ckpt=True) print(f'Resuming training from step {self.global_step}') def train(self): self.net.train() if self.opts.progressive_steps: self.check_for_progressive_training_update() while self.global_step < self.opts.max_steps: for batch_idx, batch in enumerate(self.train_dataloader): loss_dict = {} if self.is_training_discriminator(): loss_dict = self.train_discriminator(batch) x, y, y_hat, latent = self.forward(batch) loss, encoder_loss_dict, id_logs = self.calc_loss(x, y, y_hat, latent) loss_dict = {**loss_dict, **encoder_loss_dict} self.optimizer.zero_grad() loss.backward() self.optimizer.step() # Logging related if self.global_step % self.opts.image_interval == 0 or ( self.global_step < 1000 and self.global_step % 25 == 0): self.parse_and_log_images(id_logs, x, y, y_hat, title='images/train/faces') if self.global_step % self.opts.board_interval == 0: self.print_metrics(loss_dict, prefix='train') self.log_metrics(loss_dict, prefix='train') # Validation related val_loss_dict = None if self.global_step % self.opts.val_interval == 0 or self.global_step == self.opts.max_steps: val_loss_dict = self.validate() if val_loss_dict and (self.best_val_loss is None or val_loss_dict['loss'] < self.best_val_loss): self.best_val_loss = val_loss_dict['loss'] self.checkpoint_me(val_loss_dict, is_best=True) if self.global_step % self.opts.save_interval == 0 or self.global_step == self.opts.max_steps: if val_loss_dict is not None: self.checkpoint_me(val_loss_dict, is_best=False) else: self.checkpoint_me(loss_dict, is_best=False) if self.global_step == self.opts.max_steps: print('OMG, finished training!') break self.global_step += 1 if self.opts.progressive_steps: self.check_for_progressive_training_update() def check_for_progressive_training_update(self, is_resume_from_ckpt=False): for i in range(len(self.opts.progressive_steps)): if is_resume_from_ckpt and self.global_step >= self.opts.progressive_steps[i]: # Case checkpoint self.net.encoder.set_progressive_stage(ProgressiveStage(i)) if self.global_step == self.opts.progressive_steps[i]: # Case training reached progressive step self.net.encoder.set_progressive_stage(ProgressiveStage(i)) def validate(self): self.net.eval() agg_loss_dict = [] for batch_idx, batch in enumerate(self.test_dataloader): cur_loss_dict = {} if self.is_training_discriminator(): cur_loss_dict = self.validate_discriminator(batch) with torch.no_grad(): x, y, y_hat, latent = self.forward(batch) loss, cur_encoder_loss_dict, id_logs = self.calc_loss(x, y, y_hat, latent) cur_loss_dict = {**cur_loss_dict, **cur_encoder_loss_dict} agg_loss_dict.append(cur_loss_dict) # Logging related self.parse_and_log_images(id_logs, x, y, y_hat, title='images/test/faces', subscript='{:04d}'.format(batch_idx)) # For first step just do sanity test on small amount of data if self.global_step == 0 and batch_idx >= 4: self.net.train() return None # Do not log, inaccurate in first batch loss_dict = train_utils.aggregate_loss_dict(agg_loss_dict) self.log_metrics(loss_dict, prefix='test') self.print_metrics(loss_dict, prefix='test') self.net.train() return loss_dict def checkpoint_me(self, loss_dict, is_best): save_name = 'best_model.pt' if is_best else 'iteration_{}.pt'.format(self.global_step) save_dict = self.__get_save_dict() checkpoint_path = os.path.join(self.checkpoint_dir, save_name) torch.save(save_dict, checkpoint_path) with open(os.path.join(self.checkpoint_dir, 'timestamp.txt'), 'a') as f: if is_best: f.write( '**Best**: Step - {}, Loss - {:.3f} \n{}\n'.format(self.global_step, self.best_val_loss, loss_dict)) else: f.write('Step - {}, \n{}\n'.format(self.global_step, loss_dict)) def configure_optimizers(self): params = list(self.net.encoder.parameters()) if self.opts.train_decoder: params += list(self.net.decoder.parameters()) else: self.requires_grad(self.net.decoder, False) if self.opts.optim_name == 'adam': optimizer = torch.optim.Adam(params, lr=self.opts.learning_rate) else: optimizer = Ranger(params, lr=self.opts.learning_rate) return optimizer def configure_datasets(self): if self.opts.dataset_type not in data_configs.DATASETS.keys(): Exception('{} is not a valid dataset_type'.format(self.opts.dataset_type)) print('Loading dataset for {}'.format(self.opts.dataset_type)) dataset_args = data_configs.DATASETS[self.opts.dataset_type] transforms_dict = dataset_args['transforms'](self.opts).get_transforms() train_dataset = ImagesDataset(source_root=dataset_args['train_source_root'], target_root=dataset_args['train_target_root'], source_transform=transforms_dict['transform_source'], target_transform=transforms_dict['transform_gt_train'], opts=self.opts) test_dataset = ImagesDataset(source_root=dataset_args['test_source_root'], target_root=dataset_args['test_target_root'], source_transform=transforms_dict['transform_source'], target_transform=transforms_dict['transform_test'], opts=self.opts) print("Number of training samples: {}".format(len(train_dataset))) print("Number of test samples: {}".format(len(test_dataset))) return train_dataset, test_dataset def calc_loss(self, x, y, y_hat, latent): loss_dict = {} loss = 0.0 id_logs = None if self.is_training_discriminator(): # Adversarial loss loss_disc = 0. dims_to_discriminate = self.get_dims_to_discriminate() if self.is_progressive_training() else \ list(range(self.net.decoder.n_latent)) for i in dims_to_discriminate: w = latent[:, i, :] fake_pred = self.discriminator(w) loss_disc += F.softplus(-fake_pred).mean() loss_disc /= len(dims_to_discriminate) loss_dict['encoder_discriminator_loss'] = float(loss_disc) loss += self.opts.w_discriminator_lambda * loss_disc if self.opts.progressive_steps and self.net.encoder.progressive_stage.value != 18: # delta regularization loss total_delta_loss = 0 deltas_latent_dims = self.net.encoder.get_deltas_starting_dimensions() first_w = latent[:, 0, :] for i in range(1, self.net.encoder.progressive_stage.value + 1): curr_dim = deltas_latent_dims[i] delta = latent[:, curr_dim, :] - first_w delta_loss = torch.norm(delta, self.opts.delta_norm, dim=1).mean() loss_dict[f"delta{i}_loss"] = float(delta_loss) total_delta_loss += delta_loss loss_dict['total_delta_loss'] = float(total_delta_loss) loss += self.opts.delta_norm_lambda * total_delta_loss if self.opts.id_lambda > 0: # Similarity loss loss_id, sim_improvement, id_logs = self.id_loss(y_hat, y, x) loss_dict['loss_id'] = float(loss_id) loss_dict['id_improve'] = float(sim_improvement) loss += loss_id * self.opts.id_lambda if self.opts.l2_lambda > 0: loss_l2 = F.mse_loss(y_hat, y) loss_dict['loss_l2'] = float(loss_l2) loss += loss_l2 * self.opts.l2_lambda if self.opts.lpips_lambda > 0: loss_lpips = self.lpips_loss(y_hat, y) loss_dict['loss_lpips'] = float(loss_lpips) loss += loss_lpips * self.opts.lpips_lambda loss_dict['loss'] = float(loss) return loss, loss_dict, id_logs def forward(self, batch): x, y = batch x, y = x.to(self.device).float(), y.to(self.device).float() y_hat, latent = self.net.forward(x, return_latents=True) if self.opts.dataset_type == "cars_encode": y_hat = y_hat[:, :, 32:224, :] return x, y, y_hat, latent def log_metrics(self, metrics_dict, prefix): for key, value in metrics_dict.items(): self.logger.add_scalar('{}/{}'.format(prefix, key), value, self.global_step) def print_metrics(self, metrics_dict, prefix): print('Metrics for {}, step {}'.format(prefix, self.global_step)) for key, value in metrics_dict.items(): print('\t{} = '.format(key), value) def parse_and_log_images(self, id_logs, x, y, y_hat, title, subscript=None, display_count=2): im_data = [] for i in range(display_count): cur_im_data = { 'input_face': common.log_input_image(x[i], self.opts), 'target_face': common.tensor2im(y[i]), 'output_face': common.tensor2im(y_hat[i]), } if id_logs is not None: for key in id_logs[i]: cur_im_data[key] = id_logs[i][key] im_data.append(cur_im_data) self.log_images(title, im_data=im_data, subscript=subscript) def log_images(self, name, im_data, subscript=None, log_latest=False): fig = common.vis_faces(im_data) step = self.global_step if log_latest: step = 0 if subscript: path = os.path.join(self.logger.log_dir, name, '{}_{:04d}.jpg'.format(subscript, step)) else: path = os.path.join(self.logger.log_dir, name, '{:04d}.jpg'.format(step)) os.makedirs(os.path.dirname(path), exist_ok=True) fig.savefig(path) plt.close(fig) def __get_save_dict(self): save_dict = { 'state_dict': self.net.state_dict(), 'opts': vars(self.opts) } # save the latent avg in state_dict for inference if truncation of w was used during training if self.opts.start_from_latent_avg: save_dict['latent_avg'] = self.net.latent_avg if self.opts.save_training_data: # Save necessary information to enable training continuation from checkpoint save_dict['global_step'] = self.global_step save_dict['optimizer'] = self.optimizer.state_dict() save_dict['best_val_loss'] = self.best_val_loss if self.opts.w_discriminator_lambda > 0: save_dict['discriminator_state_dict'] = self.discriminator.state_dict() save_dict['discriminator_optimizer_state_dict'] = self.discriminator_optimizer.state_dict() return save_dict def get_dims_to_discriminate(self): deltas_starting_dimensions = self.net.encoder.get_deltas_starting_dimensions() return deltas_starting_dimensions[:self.net.encoder.progressive_stage.value + 1] def is_progressive_training(self): return self.opts.progressive_steps is not None # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Discriminator ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # def is_training_discriminator(self): return self.opts.w_discriminator_lambda > 0 @staticmethod def discriminator_loss(real_pred, fake_pred, loss_dict): real_loss = F.softplus(-real_pred).mean() fake_loss = F.softplus(fake_pred).mean() loss_dict['d_real_loss'] = float(real_loss) loss_dict['d_fake_loss'] = float(fake_loss) return real_loss + fake_loss @staticmethod def discriminator_r1_loss(real_pred, real_w): grad_real, = autograd.grad( outputs=real_pred.sum(), inputs=real_w, create_graph=True ) grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean() return grad_penalty @staticmethod def requires_grad(model, flag=True): for p in model.parameters(): p.requires_grad = flag def train_discriminator(self, batch): loss_dict = {} x, _ = batch x = x.to(self.device).float() self.requires_grad(self.discriminator, True) with torch.no_grad(): real_w, fake_w = self.sample_real_and_fake_latents(x) real_pred = self.discriminator(real_w) fake_pred = self.discriminator(fake_w) loss = self.discriminator_loss(real_pred, fake_pred, loss_dict) loss_dict['discriminator_loss'] = float(loss) self.discriminator_optimizer.zero_grad() loss.backward() self.discriminator_optimizer.step() # r1 regularization d_regularize = self.global_step % self.opts.d_reg_every == 0 if d_regularize: real_w = real_w.detach() real_w.requires_grad = True real_pred = self.discriminator(real_w) r1_loss = self.discriminator_r1_loss(real_pred, real_w) self.discriminator.zero_grad() r1_final_loss = self.opts.r1 / 2 * r1_loss * self.opts.d_reg_every + 0 * real_pred[0] r1_final_loss.backward() self.discriminator_optimizer.step() loss_dict['discriminator_r1_loss'] = float(r1_final_loss) # Reset to previous state self.requires_grad(self.discriminator, False) return loss_dict def validate_discriminator(self, test_batch): with torch.no_grad(): loss_dict = {} x, _ = test_batch x = x.to(self.device).float() real_w, fake_w = self.sample_real_and_fake_latents(x) real_pred = self.discriminator(real_w) fake_pred = self.discriminator(fake_w) loss = self.discriminator_loss(real_pred, fake_pred, loss_dict) loss_dict['discriminator_loss'] = float(loss) return loss_dict def sample_real_and_fake_latents(self, x): sample_z = torch.randn(self.opts.batch_size, 512, device=self.device) real_w = self.net.decoder.get_latent(sample_z) fake_w = self.net.encoder(x) if self.opts.start_from_latent_avg: fake_w = fake_w + self.net.latent_avg.repeat(fake_w.shape[0], 1, 1) if self.is_progressive_training(): # When progressive training, feed only unique w's dims_to_discriminate = self.get_dims_to_discriminate() fake_w = fake_w[:, dims_to_discriminate, :] if self.opts.use_w_pool: real_w = self.real_w_pool.query(real_w) fake_w = self.fake_w_pool.query(fake_w) if fake_w.ndim == 3: fake_w = fake_w[:, 0, :] return real_w, fake_w