def main(): args = get_args() # CUDA setting if not torch.cuda.is_available(): raise ValueError("Should buy GPU!") torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) device = torch.device('cuda') torch.set_default_tensor_type('torch.cuda.FloatTensor') torch.backends.cudnn.benchmark = True def _rescale(img): return img * 2.0 - 1.0 def _noise_adder(img): return torch.empty_like(img, dtype=img.dtype).uniform_(0.0, 1 / 128.0) + img # dataset train_dataset = datasets.ImageFolder( os.path.join(args.data_root, 'train'), transforms.Compose([ transforms.ToTensor(), _rescale, _noise_adder, ])) train_loader = iter( data.DataLoader(train_dataset, args.batch_size, sampler=InfiniteSamplerWrapper(train_dataset), num_workers=args.num_workers, pin_memory=True)) if args.calc_FID: eval_dataset = datasets.ImageFolder( os.path.join(args.data_root, 'val'), transforms.Compose([ transforms.ToTensor(), _rescale, ])) eval_loader = iter( data.DataLoader(eval_dataset, args.batch_size, sampler=InfiniteSamplerWrapper(eval_dataset), num_workers=args.num_workers, pin_memory=True)) else: eval_loader = None num_classes = len(train_dataset.classes) print(' prepared datasets...') print(' Number of training images: {}'.format(len(train_dataset))) # Prepare directories. args.num_classes = num_classes args, writer = prepare_results_dir(args) # initialize models. _n_cls = num_classes if args.cGAN else 0 gen = ResNetGenerator(args.gen_num_features, args.gen_dim_z, args.gen_bottom_width, activation=F.relu, num_classes=_n_cls, distribution=args.gen_distribution).to(device) if args.dis_arch_concat: dis = SNResNetConcatDiscriminator(args.dis_num_features, _n_cls, F.relu, args.dis_emb).to(device) else: dis = SNResNetProjectionDiscriminator(args.dis_num_features, _n_cls, F.relu).to(device) inception_model = inception.InceptionV3().to( device) if args.calc_FID else None opt_gen = optim.Adam(gen.parameters(), args.lr, (args.beta1, args.beta2)) opt_dis = optim.Adam(dis.parameters(), args.lr, (args.beta1, args.beta2)) # gen_criterion = getattr(L, 'gen_{}'.format(args.loss_type)) # dis_criterion = getattr(L, 'dis_{}'.format(args.loss_type)) gen_criterion = L.GenLoss(args.loss_type, args.relativistic_loss) dis_criterion = L.DisLoss(args.loss_type, args.relativistic_loss) print(' Initialized models...\n') if args.args_path is not None: print(' Load weights...\n') prev_args, gen, opt_gen, dis, opt_dis = utils.resume_from_args( args.args_path, args.gen_ckpt_path, args.dis_ckpt_path) # Training loop for n_iter in tqdm.tqdm(range(1, args.max_iteration + 1)): if n_iter >= args.lr_decay_start: decay_lr(opt_gen, args.max_iteration, args.lr_decay_start, args.lr) decay_lr(opt_dis, args.max_iteration, args.lr_decay_start, args.lr) # ==================== Beginning of 1 iteration. ==================== _l_g = .0 cumulative_loss_dis = .0 for i in range(args.n_dis): if i == 0: fake, pseudo_y, _ = sample_from_gen(args, device, num_classes, gen) dis_fake = dis(fake, pseudo_y) if args.relativistic_loss: real, y = sample_from_data(args, device, train_loader) dis_real = dis(real, y) else: dis_real = None loss_gen = gen_criterion(dis_fake, dis_real) gen.zero_grad() loss_gen.backward() opt_gen.step() _l_g += loss_gen.item() if n_iter % 10 == 0 and writer is not None: writer.add_scalar('gen', _l_g, n_iter) fake, pseudo_y, _ = sample_from_gen(args, device, num_classes, gen) real, y = sample_from_data(args, device, train_loader) dis_fake, dis_real = dis(fake, pseudo_y), dis(real, y) loss_dis = dis_criterion(dis_fake, dis_real) dis.zero_grad() loss_dis.backward() opt_dis.step() cumulative_loss_dis += loss_dis.item() if n_iter % 10 == 0 and i == args.n_dis - 1 and writer is not None: cumulative_loss_dis /= args.n_dis writer.add_scalar('dis', cumulative_loss_dis / args.n_dis, n_iter) # ==================== End of 1 iteration. ==================== if n_iter % args.log_interval == 0: tqdm.tqdm.write( 'iteration: {:07d}/{:07d}, loss gen: {:05f}, loss dis {:05f}'. format(n_iter, args.max_iteration, _l_g, cumulative_loss_dis)) if not args.no_image: writer.add_image( 'fake', torchvision.utils.make_grid(fake, nrow=4, normalize=True, scale_each=True)) writer.add_image( 'real', torchvision.utils.make_grid(real, nrow=4, normalize=True, scale_each=True)) # Save previews utils.save_images(n_iter, n_iter // args.checkpoint_interval, args.results_root, args.train_image_root, fake, real) if n_iter % args.checkpoint_interval == 0: # Save checkpoints! utils.save_checkpoints(args, n_iter, n_iter // args.checkpoint_interval, gen, opt_gen, dis, opt_dis) if n_iter % args.eval_interval == 0: # TODO (crcrpar): implement Ineption score, FID, and Geometry score # Once these criterion are prepared, val_loader will be used. fid_score = evaluation.evaluate(args, n_iter, gen, device, inception_model, eval_loader) tqdm.tqdm.write( '[Eval] iteration: {:07d}/{:07d}, FID: {:07f}'.format( n_iter, args.max_iteration, fid_score)) if writer is not None: writer.add_scalar("FID", fid_score, n_iter) # Project embedding weights if exists. embedding_layer = getattr(dis, 'l_y', None) if embedding_layer is not None: writer.add_embedding(embedding_layer.weight.data, list(range(args.num_classes)), global_step=n_iter) if args.test: shutil.rmtree(args.results_root)
def main(): logger = logging.getLogger('tl') args = get_args() # CUDA setting if not torch.cuda.is_available(): raise ValueError("Should buy GPU!") torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) device = torch.device('cuda') torch.set_default_tensor_type('torch.cuda.FloatTensor') torch.backends.cudnn.benchmark = True def _rescale(img): return img * 2.0 - 1.0 def _noise_adder(img): return torch.empty_like(img, dtype=img.dtype).uniform_(0.0, 1 / 128.0) + img # dataset dataset_cfg = global_cfg.get('dataset_cfg', {}) dataset_module = dataset_cfg.get('dataset_module', 'datasets.ImageFolder') train_dataset = eval(dataset_module)(os.path.join(args.data_root), transform=transforms.Compose([ transforms.ToTensor(), _rescale, _noise_adder, ]), **dataset_cfg.get( 'dataset_kwargs', {})) train_loader = iter( data.DataLoader(train_dataset, args.batch_size, sampler=InfiniteSamplerWrapper(train_dataset), num_workers=args.num_workers, pin_memory=False)) if args.calc_FID: eval_dataset = datasets.ImageFolder( os.path.join(args.data_root, 'val'), transforms.Compose([ transforms.ToTensor(), _rescale, ])) eval_loader = iter( data.DataLoader(eval_dataset, args.batch_size, sampler=InfiniteSamplerWrapper(eval_dataset), num_workers=args.num_workers, pin_memory=True)) else: eval_loader = None num_classes = len(train_dataset.classes) print(' prepared datasets...') print(' Number of training images: {}'.format(len(train_dataset))) # Prepare directories. args.num_classes = num_classes args, writer = prepare_results_dir(args) # initialize models. _n_cls = num_classes if args.cGAN else 0 gen_module = getattr( global_cfg.generator, 'module', 'pytorch_sngan_projection_lib.models.generators.resnet64') model_module = importlib.import_module(gen_module) gen = model_module.ResNetGenerator( args.gen_num_features, args.gen_dim_z, args.gen_bottom_width, activation=F.relu, num_classes=_n_cls, distribution=args.gen_distribution).to(device) if args.dis_arch_concat: dis = SNResNetConcatDiscriminator(args.dis_num_features, _n_cls, F.relu, args.dis_emb).to(device) else: dis = SNResNetProjectionDiscriminator(args.dis_num_features, _n_cls, F.relu).to(device) inception_model = inception.InceptionV3().to( device) if args.calc_FID else None opt_gen = optim.Adam(gen.parameters(), args.lr, (args.beta1, args.beta2)) opt_dis = optim.Adam(dis.parameters(), args.lr, (args.beta1, args.beta2)) # gen_criterion = getattr(L, 'gen_{}'.format(args.loss_type)) # dis_criterion = getattr(L, 'dis_{}'.format(args.loss_type)) gen_criterion = L.GenLoss(args.loss_type, args.relativistic_loss) dis_criterion = L.DisLoss(args.loss_type, args.relativistic_loss) print(' Initialized models...\n') if args.args_path is not None: print(' Load weights...\n') prev_args, gen, opt_gen, dis, opt_dis = utils.resume_from_args( args.args_path, args.gen_ckpt_path, args.dis_ckpt_path) # tf FID tf_FID = build_GAN_metric(cfg=global_cfg.GAN_metric) class SampleFunc(object): def __init__(self, generator, batch, latent, gen_distribution, device): self.generator = generator self.batch = batch self.latent = latent self.gen_distribution = gen_distribution self.device = device pass def __call__(self, *args, **kwargs): with torch.no_grad(): self.generator.eval() z = utils.sample_z(self.batch, self.latent, self.device, self.gen_distribution) pseudo_y = utils.sample_pseudo_labels(num_classes, self.batch, self.device) fake_img = self.generator(z, pseudo_y) return fake_img sample_func = SampleFunc(gen, batch=args.batch_size, latent=args.gen_dim_z, gen_distribution=args.gen_distribution, device=device) # Training loop for n_iter in tqdm.tqdm(range(1, args.max_iteration + 1)): if n_iter >= args.lr_decay_start: decay_lr(opt_gen, args.max_iteration, args.lr_decay_start, args.lr) decay_lr(opt_dis, args.max_iteration, args.lr_decay_start, args.lr) # ==================== Beginning of 1 iteration. ==================== _l_g = .0 cumulative_loss_dis = .0 for i in range(args.n_dis): if i == 0: fake, pseudo_y, _ = sample_from_gen(args, device, num_classes, gen) dis_fake = dis(fake, pseudo_y) if args.relativistic_loss: real, y = sample_from_data(args, device, train_loader) dis_real = dis(real, y) else: dis_real = None loss_gen = gen_criterion(dis_fake, dis_real) gen.zero_grad() loss_gen.backward() opt_gen.step() _l_g += loss_gen.item() if n_iter % 10 == 0 and writer is not None: writer.add_scalar('gen', _l_g, n_iter) fake, pseudo_y, _ = sample_from_gen(args, device, num_classes, gen) real, y = sample_from_data(args, device, train_loader) dis_fake, dis_real = dis(fake, pseudo_y), dis(real, y) loss_dis = dis_criterion(dis_fake, dis_real) dis.zero_grad() loss_dis.backward() opt_dis.step() cumulative_loss_dis += loss_dis.item() if n_iter % 10 == 0 and i == args.n_dis - 1 and writer is not None: cumulative_loss_dis /= args.n_dis writer.add_scalar('dis', cumulative_loss_dis / args.n_dis, n_iter) # ==================== End of 1 iteration. ==================== if n_iter % args.log_interval == 0 or n_iter == 1: tqdm.tqdm.write( 'iteration: {:07d}/{:07d}, loss gen: {:05f}, loss dis {:05f}'. format(n_iter, args.max_iteration, _l_g, cumulative_loss_dis)) if not args.no_image: writer.add_image( 'fake', torchvision.utils.make_grid(fake, nrow=4, normalize=True, scale_each=True)) writer.add_image( 'real', torchvision.utils.make_grid(real, nrow=4, normalize=True, scale_each=True)) # Save previews utils.save_images(n_iter, n_iter // args.checkpoint_interval, args.results_root, args.train_image_root, fake, real) if n_iter % args.checkpoint_interval == 0: # Save checkpoints! utils.save_checkpoints(args, n_iter, n_iter // args.checkpoint_interval, gen, opt_gen, dis, opt_dis) if (n_iter % args.eval_interval == 0 or n_iter == 1) and eval_loader is not None: # TODO (crcrpar): implement Ineption score, FID, and Geometry score # Once these criterion are prepared, val_loader will be used. fid_score = evaluation.evaluate(args, n_iter, gen, device, inception_model, eval_loader) tqdm.tqdm.write( '[Eval] iteration: {:07d}/{:07d}, FID: {:07f}'.format( n_iter, args.max_iteration, fid_score)) if writer is not None: writer.add_scalar("FID", fid_score, n_iter) # Project embedding weights if exists. embedding_layer = getattr(dis, 'l_y', None) if embedding_layer is not None: writer.add_embedding(embedding_layer.weight.data, list(range(args.num_classes)), global_step=n_iter) if n_iter % global_cfg.eval_FID_every == 0 or n_iter == 1: FID_tf, IS_mean_tf, IS_std_tf = tf_FID(sample_func=sample_func) logger.info( f'IS_mean_tf:{IS_mean_tf:.3f} +- {IS_std_tf:.3f}\n\tFID_tf: {FID_tf:.3f}' ) if not math.isnan(IS_mean_tf): summary_d = {} summary_d['FID_tf'] = FID_tf summary_d['IS_mean_tf'] = IS_mean_tf summary_d['IS_std_tf'] = IS_std_tf summary_dict2txtfig(summary_d, prefix='train', step=n_iter, textlogger=global_textlogger) gen.train() if args.test: shutil.rmtree(args.results_root)
def main(): args = get_args() # CUDA setting if not torch.cuda.is_available(): raise ValueError("Should buy GPU!") torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) device = torch.device('cuda') torch.set_default_tensor_type('torch.cuda.FloatTensor') torch.backends.cudnn.benchmark = True def _rescale(img): return img * 2.0 - 1.0 def _noise_adder(img): return torch.empty_like(img, dtype=img.dtype).uniform_(0.0, 1 / 128.0) + img eval_dataset = cifar10.CIFAR10(root=args.data_root, train=False, download=True, transform=transforms.Compose([ transforms.Resize(64), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]), minority_classes=None, keep_ratio=None) eval_loader = iter( torch.utils.data.DataLoader( eval_dataset, batch_size=args.batch_size, sampler=InfiniteSamplerWrapper(eval_dataset), num_workers=args.num_workers, pin_memory=True)) print(' prepared datasets...') # Prepare directories. num_classes = len(eval_dataset.classes) args.num_classes = num_classes # initialize models. _n_cls = num_classes if args.cGAN else 0 gen = ResNetGenerator(args.gen_num_features, args.gen_dim_z, args.gen_bottom_width, activation=F.relu, num_classes=_n_cls, distribution=args.gen_distribution).to(device) if args.dis_arch_concat: dis = SNResNetConcatDiscriminator(args.dis_num_features, _n_cls, F.relu, args.dis_emb).to(device) else: dis = SNResNetProjectionDiscriminator(args.dis_num_features, _n_cls, F.relu, args.transform_space).to(device) inception_model = inception.InceptionV3().to( device) if args.calc_FID else None gen = torch.nn.DataParallel(gen) # dis = torch.nn.DataParallel(dis) opt_gen = optim.Adam(gen.parameters(), args.lr, (args.beta1, args.beta2)) opt_dis = optim.Adam(dis.parameters(), args.lr, (args.beta1, args.beta2)) # gen_criterion = getattr(L, 'gen_{}'.format(args.loss_type)) # dis_criterion = getattr(L, 'dis_{}'.format(args.loss_type)) gen_criterion = L.GenLoss(args.loss_type, args.relativistic_loss) dis_criterion = L.DisLoss(args.loss_type, args.relativistic_loss) print(' Initialized models...\n') if args.args_path is None: print("Please specify weights to load") exit() else: print(' Load weights...\n') prev_args, gen, opt_gen, dis, opt_dis = utils.resume_from_args( args.args_path, args.gen_ckpt_path, args.dis_ckpt_path) args.n_fid_batches = args.n_eval_batches fid_score = evaluation.evaluate(args, 0, gen, device, inception_model, eval_loader, to_save=False) print(fid_score)