From 17555e7d075c72f9175f4e03d33c32856ed78673 Mon Sep 17 00:00:00 2001 From: James Betker Date: Wed, 21 Apr 2021 18:14:17 -0600 Subject: [PATCH] misc adjustments for stylegan --- codes/models/stylegan/stylegan2_rosinality.py | 7 ++++++- codes/scripts/byol/byol_resnet_playground.py | 8 +++++--- codes/test.py | 9 ++++++--- codes/train.py | 2 +- codes/trainer/ExtensibleTrainer.py | 4 ++-- codes/trainer/eval/fid.py | 4 ++++ 6 files changed, 24 insertions(+), 10 deletions(-) diff --git a/codes/models/stylegan/stylegan2_rosinality.py b/codes/models/stylegan/stylegan2_rosinality.py index e97041dc..90a14863 100644 --- a/codes/models/stylegan/stylegan2_rosinality.py +++ b/codes/models/stylegan/stylegan2_rosinality.py @@ -558,7 +558,12 @@ class Generator(nn.Module): randomize_noise=True, ): if not input_is_latent: - styles = [self.style(s) for s in styles] + if self.training: + # In training mode, multiple style vectors are fed to the generator. + styles = [self.style(s) for s in styles] + else: + # In eval mode, only a single style is fed. + styles = [self.style(styles)] if noise is None: if randomize_noise: diff --git a/codes/scripts/byol/byol_resnet_playground.py b/codes/scripts/byol/byol_resnet_playground.py index cd74d36b..2fd49975 100644 --- a/codes/scripts/byol/byol_resnet_playground.py +++ b/codes/scripts/byol/byol_resnet_playground.py @@ -57,8 +57,9 @@ def get_image_folder_dataloader(batch_size, num_workers, target_size=224, shuffl 'name': 'amalgam', #'paths': ['F:\\4k6k\\datasets\\ns_images\\imagesets\\pn_coven\\cropped2'], #'paths': ['F:\\4k6k\\datasets\\ns_images\\imagesets\\imageset_1024_square_with_new'], - 'paths': ['F:\\4k6k\\datasets\\ns_images\\imagesets\\imageset_256_tiled_filtered_flattened'], + #'paths': ['F:\\4k6k\\datasets\\ns_images\\imagesets\\imageset_256_tiled_filtered_flattened'], #'paths': ['F:\\4k6k\\datasets\\ns_images\\imagesets\\1024_test'], + 'paths': ['E:\\4k6k\\datasets\\ns_images\\imagesets\\imageset_256_full'], 'weights': [1], 'target_size': target_size, 'force_multiple': 32, @@ -116,6 +117,7 @@ def produce_latent_dict(model): latents = [] for batch in tqdm(dataloader): hq = batch['hq'].to('cuda') + hq = F.interpolate(F.interpolate(hq, size=(16,16), mode='bilinear'), size=(224,244)) model(hq) l = layer_hooked_value.cpu().split(1, dim=0) latents.extend(l) @@ -202,7 +204,7 @@ if __name__ == '__main__': register_hook(model, 'avgpool') with torch.no_grad(): - find_similar_latents(model, structural_euc_dist) - #produce_latent_dict(model) + #find_similar_latents(model, structural_euc_dist) + produce_latent_dict(model) #build_kmeans() #use_kmeans() diff --git a/codes/test.py b/codes/test.py index b3d4b4bc..6a2529f8 100644 --- a/codes/test.py +++ b/codes/test.py @@ -13,11 +13,14 @@ from tqdm import tqdm import torch -def forward_pass(model, output_dir, alteration_suffix=''): +def forward_pass(model, output_dir, opt): + alteration_suffix = util.opt_get(opt, ['name'], '') + denorm_range = tuple(util.opt_get(opt, ['image_normalization_range'], [0, 1])) model.feed_data(data, 0, need_GT=need_GT) model.test() visuals = model.get_current_visuals(need_GT)['rlt'].cpu() + visuals = (visuals - denorm_range[0]) / (denorm_range[1]-denorm_range[0]) fea_loss = 0 psnr_loss = 0 for i in range(visuals.shape[0]): @@ -48,7 +51,7 @@ if __name__ == "__main__": torch.backends.cudnn.benchmark = True want_metrics = False parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to options YAML file.', default='../options/test_mi1.yml') + parser.add_argument('-opt', type=str, help='Path to options YAML file.', default='../options/test_cats_stylegan2_rosinality.yml') opt = option.parse(parser.parse_args().opt, is_train=False) opt = option.dict_to_nonedict(opt) utils.util.loaded_options = opt @@ -90,7 +93,7 @@ if __name__ == "__main__": need_GT = False if test_loader.dataset.opt['dataroot_GT'] is None else True need_GT = need_GT and want_metrics - fea_loss, psnr_loss = forward_pass(model, dataset_dir, opt['name']) + fea_loss, psnr_loss = forward_pass(model, dataset_dir, opt) fea_loss += fea_loss psnr_loss += psnr_loss diff --git a/codes/train.py b/codes/train.py index 26329b50..2fdfb4b3 100644 --- a/codes/train.py +++ b/codes/train.py @@ -295,7 +295,7 @@ class Trainer: if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_lightweight_gan_pna.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_cats_stylegan2_rosinality.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() diff --git a/codes/trainer/ExtensibleTrainer.py b/codes/trainer/ExtensibleTrainer.py index f137cf83..946258e9 100644 --- a/codes/trainer/ExtensibleTrainer.py +++ b/codes/trainer/ExtensibleTrainer.py @@ -259,8 +259,8 @@ class ExtensibleTrainer(BaseModel): # Record visual outputs for usage in debugging and testing. if 'visuals' in self.opt['logger'].keys() and self.rank <= 0 and step % self.opt['logger']['visual_debug_rate'] == 0: - denorm = opt_get(self.opt, ['logger', 'denormalize'], False) - denorm_range = opt_get(self.opt, ['logger', 'denormalize_range'], None) + denorm = 'image_normalization_range' in self.opt.keys() + denorm_range = opt_get(self.opt, ['image_normalization_range'], None) if denorm_range: denorm_range = tuple(denorm_range) sample_save_path = os.path.join(self.opt['path']['models'], "..", "visual_dbg") diff --git a/codes/trainer/eval/fid.py b/codes/trainer/eval/fid.py index eb591a5f..f125fae7 100644 --- a/codes/trainer/eval/fid.py +++ b/codes/trainer/eval/fid.py @@ -18,11 +18,13 @@ class StyleTransferEvaluator(evaluator.Evaluator): self.gen_output_index = opt_eval['gen_index'] if 'gen_index' in opt_eval.keys() else 0 self.noise_type = opt_get(opt_eval, ['noise_type'], 'imgnoise') self.latent_dim = opt_get(opt_eval, ['latent_dim'], 512) # Not needed if using 'imgnoise' input. + self.image_norm_range = tuple(opt_get(env['opt'], ['image_normalization_range'], [0,1])) def perform_eval(self): fid_fake_path = osp.join(self.env['base_path'], "../", "fid", str(self.env["step"])) os.makedirs(fid_fake_path, exist_ok=True) counter = 0 + self.model.eval() for i in range(self.batches_per_eval): if self.noise_type == 'imgnoise': batch = torch.FloatTensor(self.batch_sz, 3, self.im_sz, self.im_sz).uniform_(0., 1.).to(self.env['device']) @@ -32,9 +34,11 @@ class StyleTransferEvaluator(evaluator.Evaluator): if not isinstance(gen, list) and not isinstance(gen, tuple): gen = [gen] gen = gen[self.gen_output_index] + gen = (gen - self.image_norm_range[0]) / (self.image_norm_range[1]-self.image_norm_range[0]) for b in range(self.batch_sz): torchvision.utils.save_image(gen[b], osp.join(fid_fake_path, "%i_.png" % (counter))) counter += 1 + self.model.train() print("Got all images, computing fid") return {"fid": fid_score.calculate_fid_given_paths([self.fid_real_samples, fid_fake_path], self.batch_sz, True,