def load_model(model_file, hidden_size, upsampling, cuda=False): if cuda: from_before = torch.load(model_file) else: from_before = torch.load(model_file, map_location=lambda storage, loc: storage) total_examples = from_before['total_examples'] gen_losses = from_before['gen_losses'] disc_losses = from_before['disc_losses'] gen_loss_per_epoch = from_before['gen_loss_per_epoch'] disc_loss_per_epoch = from_before['disc_loss_per_epoch'] gen_state_dict = from_before['gen_state_dict'] disc_state_dict = from_before['disc_state_dict'] fixed_noise = from_before['fixed_noise'] epoch = from_before['epoch'] # load generator and discriminator if upsampling == 'transpose': from models.model import Generator, Discriminator elif upsampling == 'nn': from models.model_nn import Generator, Discriminator elif upsampling == 'bilinear': from models.model_bilinear import Generator, Discriminator gen = Generator(hidden_dim=hidden_size, dropout=0.4) # TODO: save dropout in checkpoint disc = Discriminator(leaky=0.2, dropout=0.4) # TODO: same here disc.load_state_dict(disc_state_dict) gen.load_state_dict(gen_state_dict) return total_examples, fixed_noise, gen_losses, disc_losses, \ gen_loss_per_epoch, disc_loss_per_epoch, epoch, gen, disc
def test_shape_g(self): gen = Generator(self.noise_dim, self.in_channels, 8) self.assertEqual( gen(self.z).shape, (self.N, self.in_channels, self.H, self.W)) print( f"{gen(self.z).shape} == {(self.N, self.in_channels, self.H, self.W)}" )
def train(): args = parse_args() cfg = Config.from_file(args.config) # Dimensionality of the latent vector. latent_size = cfg.models.generator.z_dim # Use sigmoid activation for the last layer? cfg.models.discriminator.sigmoid_at_end = cfg.train.loss_type in [ 'ls', 'gan' ] os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) G = Generator(model_cfg=cfg.models.generator, target_size=cfg.train.target_size) D = Discriminator(model_cfg=cfg.models.discriminator, target_size=cfg.train.target_size) #print(G) #print(D) dataset = FaceDataset(cfg.train.dataset) assert len(dataset) > 0 print(f'train dataset contains {len(dataset)} images.') clip = cfg.models.generator.z_clipping if hasattr(cfg.models.generator, 'z_zlipping') else None z_generator = RandomNoiseGenerator(cfg.models.generator.z_dim, 'gaussian', clip=clip) pggan = PGGAN(G, D, dataset, z_generator, args.gpu, cfg, args.resume) pggan.train()
def model_init(self): self.D = Discriminator(self.model_type, self.image_size, self.hidden_dim, self.n_filter, self.n_repeat) self.G = Generator(self.model_type, self.image_size, self.hidden_dim, self.n_filter, self.n_repeat) self.D = cuda(self.D, self.cuda) self.G = cuda(self.G, self.cuda) self.D.weight_init(mean=0.0, std=0.02) self.G.weight_init(mean=0.0, std=0.02) self.D_optim = optim.Adam(self.D.parameters(), lr=self.D_lr, betas=(0.5, 0.999)) self.G_optim = optim.Adam(self.G.parameters(), lr=self.G_lr, betas=(0.5, 0.999)) #self.D_optim_scheduler = lr_scheduler.ExponentialLR(self.D_optim, gamma=0.97) #self.G_optim_scheduler = lr_scheduler.ExponentialLR(self.G_optim, gamma=0.97) self.D_optim_scheduler = lr_scheduler.StepLR(self.D_optim, step_size=1, gamma=0.5) self.G_optim_scheduler = lr_scheduler.StepLR(self.G_optim, step_size=1, gamma=0.5) if not self.ckpt_dir.exists(): self.ckpt_dir.mkdir(parents=True, exist_ok=True) if self.load_ckpt: self.load_checkpoint()
def load_model(opt): print('initializing model ...') print("LOADING GENERATOR MODEL") model_g = Generator(100, gpu=opt.SYSTEM.USE_GPU) print("LOADING DISCRIMINATOR MODEL") model_d = Discriminator(output_dim=opt.TRAIN.TOTAL_FEATURES, gpu=opt.SYSTEM.USE_GPU) return model_g, model_d
def test_generator_shape(self): gen_sle = Generator(img_size=self.img_size, in_channels=self.in_channels, z_dim=self.z_dim, res_type=self.res_type[0]) gen_gc = Generator(img_size=self.img_size, in_channels=self.in_channels, z_dim=self.z_dim, res_type=self.res_type[1]) noise = torch.randn([1, 256, 1, 1]) z_gen_sle = gen_sle(noise) z_gen_gc = gen_gc(noise) real_shape = (1, 3, self.img_size, self.img_size) self.assertEqual(z_gen_sle.shape, z_gen_gc.shape) self.assertEqual(z_gen_sle.shape, real_shape) self.assertEqual(z_gen_gc.shape, real_shape)
def test_shape(self): gen = Generator(self.Z_DIM, self.IN_CHANNELS, img_channels=3) critic = Critic(self.IN_CHANNELS, img_channels=3) for img_size in [4, 8, 16, 32, 64, 128, 256, 512, 1024]: num_steps = int(log2(img_size / 4)) x = torch.randn((1, self.Z_DIM, 1, 1)) z = gen(x, 0.5, steps=num_steps) self.assertEqual(z.shape, (1, 3, img_size, img_size)) out = critic(z, alpha=0.5, steps=num_steps) self.assertEqual(out.shape, (1, 1)) print(f"Done! Img_size:{img_size}")
def models_initialize(self): print("-- Preparing Models --") self.generator = Generator(ngf=self.args.MODEL.ngf, input_nc=self.args.MODEL.input_nc, output_nc=self.args.MODEL.output_nc) self.discriminator = NLayerDiscriminator( ndf=self.args.MODEL.ndf, input_nc=self.args.MODEL.input_nc + self.args.MODEL.output_nc, n_layers=self.args.MODEL.n_layers) self.generator.to(self.device) self.discriminator.to(self.device) print("-- Models DONE --")
def test_fid_with_images(self): gen = Generator(1024) fid_model = InceptionV3FID(torch.device('cpu')) real_dataset = ImgFolderDataset('', fid=True) real_dataloader = get_sample_dataloader(real_dataset, num_samples=4, batch_size=2) noise = torch.randn([len(real_dataloader), 256, 1, 1]) fake_images = [] for batch in noise: fake_images.append(gen(batch.unsqueeze(0))) noise_dataset = FIDNoiseDataset(fake_images) fake_dataloader = DataLoader(noise_dataset, batch_size=2) fid = fid_model.get_fid_score(real_dataloader, fake_dataloader) print(fid)
def __init__(self, args, device): self.args = args self.device = device # Load generator self.generator = Generator( args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier).to(device) if args.generator_ckpt is not None: print("load generator:", args.generator_ckpt) generator_checkpoint = torch.load( args.generator_ckpt ) # map_location=lambda storage, loc: storage) self.generator.load_state_dict( generator_checkpoint['g_ema'], strict=False) # TODO: Maybe need to load g # Load classifier resnet = resnet50(pretrained=True, num_classes=1) self.attribute_classifier = resnet.to(device) if args.attribute_classifier_ckpt is not None: print("load attribute clasifier:", args.attribute_classifier_ckpt) classifier_checkpoint = torch.load(args.attribute_classifier_ckpt) self.attribute_classifier.load_state_dict( classifier_checkpoint['state_dict']) self.optimizer = torch.optim.AdamW( self.generator.attribute_mapper.parameters(), lr=args.lr) self.classifier_criterion = nn.BCEWithLogitsLoss() self.mse_loss = nn.MSELoss() self.sample_z = None # Sample for validation self.mean_latent = None # Mean for truncation cudnn.benchmark = True self.writer = SummaryWriter(log_dir=str(args.log_dir))
if args.gif and args.num_samples < 2: raise ValueError('for GIF num_samples must be greater than 1') if not args.out_path: out_path = 'DCGAN-Anime-Faces' else: out_path = args.out_path if args.device == 'cuda': if torch.cuda.is_available(): device = torch.device('cuda') else: device = torch.device('cpu') gen = Generator(128, 3, 64) load_gen(gen, args.path_ckpt, device) gen.eval() if args.grid: noise = get_random_noise(args.num_samples, args.z_size, device) print("==> Generate IMAGE GRID...") output = gen(noise) show_batch(output, out_path, num_samples=args.num_samples, figsize=(args.img_size, args.img_size)) elif args.gif: noise = get_random_noise(args.num_samples, args.z_size, device) print("==> Generate GIF...") images = latent_space_interpolation_sequence(noise, step_interpolation=args.steps) output = gen(images) if args.resize and isinstance(args.resize, int): print(f"==> Resize images to {args.resize}px")
def main(args): # set which gpu(s) to use, should set PCI_BUS_ID first os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu num_gpus = (len(args.gpu) + 1) // 2 # create model directories checkpath(args.modelG_path) checkpath(args.modelD_path) # tensorboard writer checkpath(args.log_path) writer = SummaryWriter(args.log_path) # load data data_loader, num_train = get_loader(args, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, training=True) data_loader_val, num_test = get_loader(args, batch_size=args.val_bs, shuffle=False, num_workers=args.num_workers, training=False) print('Finished data loading') print("The length of the train set is: {}".format(num_train)) print("The length of the test set is: {}".format(num_test)) colorguide = True if args.nocolor: colorguide = False # loss multipliers lambdas = [ args.lambda_imgl1, args.lambda_wfl1, args.lambda_ssim, args.lambda_color ] lambda_perceptual = args.lambda_perceptual # Generator netG = Generator(lambdas=lambdas, colorguide=colorguide, input_nc=1, output_nc=1) if num_gpus > 1: # multi-gpu training with synchonized batchnormalization # make sure enough number of gpus are available assert (torch.cuda.device_count() >= num_gpus) # since we have set CUDA_VISIBLE_DEVICES to avoid some invalid device id issues netG = DataParallelWithCallback( netG, device_ids=[i for i in range(num_gpus)]) netG_single = netG.module else: # single gpu training netG_single = netG # Discriminator netD = NLayerDiscriminator(input_nc=4, n_layers=4) if num_gpus > 1: netD = DataParallelWithCallback( netD, device_ids=[i for i in range(num_gpus)]) netD_single = netD.module else: netD_single = netD # print(netG_single) # print(netD_single) if args.pretrained and args.netG_path != '' and args.netD_path != '': netG_single.load_state_dict(torch.load(args.netG_path)) netD_single.load_state_dict(torch.load(args.netD_path)) # Right now we only support gpu training if torch.cuda.is_available(): netG = netG.cuda() netD = netD.cuda() # define the perceptual loss, place outside the forward func in G for better multi-gpu training Ploss = PNet() if num_gpus > 1: Ploss = DataParallelWithCallback( Ploss, device_ids=[i for i in range(num_gpus)]) if torch.cuda.is_available(): Ploss = Ploss.cuda() # setup optimizer lr = args.learning_rate optimizerD = optim.Adam(netD_single.parameters(), lr=lr, betas=(args.beta1, 0.999)) schedulerD = ReduceLROnPlateau(optimizerD, factor=0.7, patience=10, mode='min', min_lr=1e-06) optimizerG = optim.Adam(netG_single.parameters(), lr=lr, betas=(args.beta1, 0.999)) schedulerG = ReduceLROnPlateau(optimizerG, factor=0.7, patience=10, mode='min', min_lr=1e-06) for epoch in range(args.num_epochs): # switch to train mode netG.train() netD.train() for i, (img_real, wf_real, color_real) in enumerate(data_loader, 0): img_real = img_real.cuda() wf_real = wf_real.cuda() color_real = color_real.cuda() # Update D network, we freeze parameters in G to save memory for p in netG_single.parameters(): p.requires_grad = False for p in netD_single.parameters(): p.requires_grad = True # if using TTUR, D can be trained multiple steps per G step for _ in range(args.D_steps): optimizerD.zero_grad() # train with real real_AB = torch.cat((img_real, wf_real), 1) errD_real = 0.5 * netD(trainG=False, trainReal=True, real_AB=real_AB, fake_AB=None).sum() errD_real.backward() # train with fake img_fake, wf_fake, _, _, _, _, _ = netG(trainG=False, img_real=None, wf_real=wf_real, color_real=color_real) fake_AB = torch.cat((img_fake, wf_fake), 1) errD_fake = 0.5 * netD(trainG=False, trainReal=False, real_AB=None, fake_AB=fake_AB).sum() errD_fake.backward() errD = errD_real + errD_fake optimizerD.step() del img_fake, wf_fake, fake_AB, real_AB, errD_real, errD_fake iterations_before_epoch = epoch * len(data_loader) writer.add_scalar('D Loss', errD.item(), iterations_before_epoch + i) del errD # Update G network, we freeze parameters in D to save memory for p in netG.parameters(): p.requires_grad = True for p in netD.parameters(): p.requires_grad = False optimizerG.zero_grad() img_fake, wf_fake, lossG, wf_ssim, img_l1, color_l1, wf_l1 = netG( trainG=True, img_real=img_real, wf_real=wf_real, color_real=color_real) ploss = Ploss(img_fake, img_real.detach()).sum() fake_AB = torch.cat((img_fake, wf_fake), 1) lossD = netD(trainG=True, trainReal=False, real_AB=None, fake_AB=fake_AB).sum() errG = (lossG.sum() + lambda_perceptual * ploss + lossD) errG.backward() optimizerG.step() del color_real, fake_AB, lossG, errG if args.nocolor: print( 'Epoch: [{}/{}] Iter: [{}/{}] PercLoss : {:.4f} ImageL1 : {:.6f} WfL1 : {:.6f} WfSSIM : {:.6f}' .format(epoch, args.num_epochs, i, len(data_loader), ploss.item(), img_l1.sum().item(), wf_l1.sum().item(), num_gpus + wf_ssim.sum().item())) else: print( 'Epoch: [{}/{}] Iter: [{}/{}] PercLoss : {:.4f} ImageL1 : {:.6f} WfL1 : {:.6f} WfSSIM : {:.6f} ColorL1 : {:.6f}' .format(epoch, args.num_epochs, i, len(data_loader), ploss.item(), img_l1.sum().item(), wf_l1.sum().item(), num_gpus + wf_ssim.sum().item(), color_l1.sum().item())) writer.add_scalar('Color Loss', color_l1.sum().item(), iterations_before_epoch + i) # tensorboard log writer.add_scalar('G Loss', lossD.item(), iterations_before_epoch + i) writer.add_scalar('Image L1 Loss', img_l1.sum().item(), iterations_before_epoch + i) writer.add_scalar('Wireframe MSSSIM Loss', num_gpus + wf_ssim.sum().item(), iterations_before_epoch + i) writer.add_scalar('Wireframe L1', wf_l1.sum().item(), iterations_before_epoch + i) writer.add_scalar('Image Perceptual Loss', ploss.item(), iterations_before_epoch + i) del wf_ssim, ploss, img_l1, color_l1, wf_l1, lossD with torch.no_grad(): # show generated tarining images in tensorboard if i % args.val_freq == 0: real_img = vutils.make_grid( img_real.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Real Image', real_img, (iterations_before_epoch + i) // args.val_freq) real_wf = vutils.make_grid( wf_real.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Real Wireframe', real_wf, (iterations_before_epoch + i) // args.val_freq) fake_img = vutils.make_grid( img_fake.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Fake Image', fake_img, (iterations_before_epoch + i) // args.val_freq) fake_wf = vutils.make_grid( wf_fake.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Fake Wireframe', fake_wf, (iterations_before_epoch + i) // args.val_freq) del real_img, real_wf, fake_img, fake_wf del img_real, wf_real, img_fake, wf_fake # do checkpointing if epoch % args.save_freq == 0 and epoch > 0: torch.save(netG_single.state_dict(), '{}/netG_epoch_{}.pth'.format(args.modelG_path, epoch)) torch.save(netD_single.state_dict(), '{}/netD_epoch_{}.pth'.format(args.modelD_path, epoch)) # validation with torch.no_grad(): netG_single.eval() # since we use a realtively large validation batchsize, we don't go through the who test set (img_real, wf_real, color_real) = next(iter(data_loader_val)) img_real = img_real.cuda() wf_real = wf_real.cuda() color_real = color_real.cuda() img_fake, wf_fake, _, _, _, _, _ = netG_single( trainG=False, img_real=None, wf_real=wf_real, color_real=color_real) # update lr based on the validation perceptual loss val_score = Ploss(img_fake.detach(), img_real.detach()).sum() schedulerG.step(val_score) schedulerD.step(val_score) print('Current lr: {:.6f}'.format( optimizerG.param_groups[0]['lr'])) real_img = vutils.make_grid(img_real.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Test: Real Image', real_img, epoch) real_wf = vutils.make_grid(wf_real.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Test: Real Wireframe', real_wf, epoch) fake_img = vutils.make_grid(img_fake.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Test: Fake Image', fake_img, epoch) fake_wf = vutils.make_grid(wf_fake.detach()[:args.val_size], normalize=True, scale_each=True) writer.add_image('Test: Fake Wireframe', fake_wf, epoch) netG_single.train() del img_real, real_img, wf_real, real_wf, img_fake, fake_img, wf_fake, fake_wf # close tb writer writer.close()
def main(train_set, learning_rate, n_epochs, beta_0, beta_1, batch_size, num_workers, hidden_size, model_file, cuda, display_result_every, checkpoint_interval, seed, label_smoothing, grad_clip, dropout, upsampling): # make data between -1 and 1 data_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) train_dataset = datasets.ImageFolder(root=os.path.join( os.getcwd(), train_set), transform=data_transform) train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True) # initialize model if model_file: try: total_examples, fixed_noise, gen_losses, disc_losses, gen_loss_per_epoch, \ disc_loss_per_epoch, prev_epoch, gen, disc = load_model(model_file, hidden_size, upsampling, cuda) print('model loaded successfully!') except: print('could not load model! creating new model...') model_file = None if not model_file: print('creating new model...') if upsampling == 'transpose': from models.model import Generator, Discriminator elif upsampling == 'nn': from models.model_nn import Generator, Discriminator elif upsampling == 'bilinear': from models.model_bilinear import Generator, Discriminator gen = Generator(hidden_dim=hidden_size, leaky=0.2, dropout=dropout) disc = Discriminator(leaky=0.2, dropout=dropout) gen.weight_init(mean=0, std=0.02) disc.weight_init(mean=0, std=0.02) total_examples = 0 disc_losses = [] gen_losses = [] disc_loss_per_epoch = [] gen_loss_per_epoch = [] prev_epoch = 0 # Sample minibatch of m noise samples from noise prior p_g(z) and transform if cuda: fixed_noise = Variable(torch.randn(9, hidden_size).cuda()) else: fixed_noise = Variable(torch.rand(9, hidden_size)) if cuda: gen.cuda() disc.cuda() # Binary Cross Entropy loss BCE_loss = nn.BCELoss() # Adam optimizer gen_optimizer = optim.Adam(gen.parameters(), lr=learning_rate, betas=(beta_0, beta_1), eps=1e-8) disc_optimizer = optim.Adam(disc.parameters(), lr=learning_rate, betas=(beta_0, beta_1), eps=1e-8) # results save folder gen_images_dir = 'results/generated_images' train_summaries_dir = 'results/training_summaries' checkpoint_dir = 'results/checkpoints' if not os.path.isdir('results'): os.mkdir('results') if not os.path.isdir(gen_images_dir): os.mkdir(gen_images_dir) if not os.path.isdir(train_summaries_dir): os.mkdir(train_summaries_dir) if not os.path.isdir(checkpoint_dir): os.mkdir(checkpoint_dir) np.random.seed( seed ) # reset training seed to ensure that batches remain the same between runs! try: for epoch in range(prev_epoch, n_epochs): disc_losses_epoch = [] gen_losses_epoch = [] for idx, (true_batch, _) in enumerate(train_dataloader): disc.zero_grad() # hack 6 of https://github.com/soumith/ganhacks if label_smoothing: true_target = torch.FloatTensor(batch_size).uniform_( 0.7, 1.2) else: true_target = torch.ones(batch_size) # Sample minibatch of examples from data generating distribution if cuda: true_batch = Variable(true_batch.cuda()) true_target = Variable(true_target.cuda()) else: true_batch = Variable(true_batch) true_target = Variable(true_target) # train discriminator on true data true_disc_result = disc.forward(true_batch) disc_train_loss_true = BCE_loss(true_disc_result.squeeze(), true_target) disc_train_loss_true.backward() torch.nn.utils.clip_grad_norm(disc.parameters(), grad_clip) # Sample minibatch of m noise samples from noise prior p_g(z) and transform if label_smoothing: fake_target = torch.FloatTensor(batch_size).uniform_( 0, 0.3) else: fake_target = torch.zeros(batch_size) if cuda: z = Variable(torch.randn(batch_size, hidden_size).cuda()) fake_target = Variable(fake_target.cuda()) else: z = Variable(torch.randn(batch_size, hidden_size)) fake_target = Variable(fake_target) # train discriminator on fake data fake_batch = gen.forward(z.view(-1, hidden_size, 1, 1)) fake_disc_result = disc.forward(fake_batch.detach( )) # detach so gradients not computed for generator disc_train_loss_false = BCE_loss(fake_disc_result.squeeze(), fake_target) disc_train_loss_false.backward() torch.nn.utils.clip_grad_norm(disc.parameters(), grad_clip) disc_optimizer.step() # compute performance statistics disc_train_loss = disc_train_loss_true + disc_train_loss_false disc_losses_epoch.append(disc_train_loss.data[0]) disc_fake_accuracy = 1 - fake_disc_result.mean().data[0] disc_true_accuracy = true_disc_result.mean().data[0] # Sample minibatch of m noise samples from noise prior p_g(z) and transform if label_smoothing: true_target = torch.FloatTensor(batch_size).uniform_( 0.7, 1.2) else: true_target = torch.ones(batch_size) if cuda: z = Variable(torch.randn(batch_size, hidden_size).cuda()) true_target = Variable(true_target.cuda()) else: z = Variable(torch.rand(batch_size, hidden_size)) true_target = Variable(true_target) # train generator gen.zero_grad() fake_batch = gen.forward(z.view(-1, hidden_size, 1, 1)) disc_result = disc.forward(fake_batch) gen_train_loss = BCE_loss(disc_result.squeeze(), true_target) gen_train_loss.backward() torch.nn.utils.clip_grad_norm(gen.parameters(), grad_clip) gen_optimizer.step() gen_losses_epoch.append(gen_train_loss.data[0]) if (total_examples != 0) and (total_examples % display_result_every == 0): print( 'epoch {}: step {}/{} disc true acc: {:.4f} disc fake acc: {:.4f} ' 'disc loss: {:.4f}, gen loss: {:.4f}'.format( epoch + 1, idx + 1, len(train_dataloader), disc_true_accuracy, disc_fake_accuracy, disc_train_loss.data[0], gen_train_loss.data[0])) # Checkpoint model total_examples += batch_size if (total_examples != 0) and (total_examples % checkpoint_interval == 0): disc_losses.extend(disc_losses_epoch) gen_losses.extend(gen_losses_epoch) save_checkpoint(total_examples=total_examples, fixed_noise=fixed_noise, disc=disc, gen=gen, gen_losses=gen_losses, disc_losses=disc_losses, disc_loss_per_epoch=disc_loss_per_epoch, gen_loss_per_epoch=gen_loss_per_epoch, epoch=epoch, directory=checkpoint_dir) print("Checkpoint saved!") # sample images for inspection save_image_sample(batch=gen.forward( fixed_noise.view(-1, hidden_size, 1, 1)), cuda=cuda, total_examples=total_examples, directory=gen_images_dir) print("Saved images!") # save learning curves for inspection save_learning_curve(gen_losses=gen_losses, disc_losses=disc_losses, total_examples=total_examples, directory=train_summaries_dir) print("Saved learning curves!") disc_loss_per_epoch.append(np.average(disc_losses_epoch)) gen_loss_per_epoch.append(np.average(gen_losses_epoch)) # Save epoch learning curve save_learning_curve_epoch(gen_losses=gen_loss_per_epoch, disc_losses=disc_loss_per_epoch, total_epochs=epoch + 1, directory=train_summaries_dir) print("Saved learning curves!") print('epoch {}/{} disc loss: {:.4f}, gen loss: {:.4f}'.format( epoch + 1, n_epochs, np.array(disc_losses_epoch).mean(), np.array(gen_losses_epoch).mean())) disc_losses.extend(disc_losses_epoch) gen_losses.extend(gen_losses_epoch) except KeyboardInterrupt: print("Saving before quit...") save_checkpoint(total_examples=total_examples, fixed_noise=fixed_noise, disc=disc, gen=gen, disc_loss_per_epoch=disc_loss_per_epoch, gen_loss_per_epoch=gen_loss_per_epoch, gen_losses=gen_losses, disc_losses=disc_losses, epoch=epoch, directory=checkpoint_dir) print("Checkpoint saved!") # sample images for inspection save_image_sample(batch=gen.forward( fixed_noise.view(-1, hidden_size, 1, 1)), cuda=cuda, total_examples=total_examples, directory=gen_images_dir) print("Saved images!") # save learning curves for inspection save_learning_curve(gen_losses=gen_losses, disc_losses=disc_losses, total_examples=total_examples, directory=train_summaries_dir) print("Saved learning curves!")
"--sampling", default="end", choices=["end", "full"], help="set endpoint sampling method", ) parser.add_argument("ckpt", metavar="CHECKPOINT", help="path to the model checkpoints") args = parser.parse_args() latent_dim = 512 ckpt = torch.load(args.ckpt) g = Generator(args.size, latent_dim, 8).to(device) g.load_state_dict(ckpt["g_ema"]) g.eval() percept = lpips.PerceptualLoss(model="net-lin", net="vgg", use_gpu=device.startswith("cuda")) distances = [] n_batch = args.n_sample // args.batch resid = args.n_sample - (n_batch * args.batch) batch_sizes = [args.batch] * n_batch + [resid] with torch.no_grad(): for batch in tqdm(batch_sizes):
help='rampup_kimg.') parser.add_argument('--rampdown_kimg', default=10000, type=float, help='rampdown_kimg.') # TODO: support conditional inputs args = parser.parse_args() opts = {k: v for k, v in args._get_kwargs()} latent_size = 512 sigmoid_at_end = args.gan in ['lsgan', 'gan'] G = Generator(num_channels=3, latent_size=latent_size, resolution=args.target_resol, fmap_max=latent_size, fmap_base=8192, tanh_at_end=False) D = Discriminator(num_channels=3, resolution=args.target_resol, fmap_max=latent_size, fmap_base=8192, sigmoid_at_end=sigmoid_at_end) print(G) print(D) data = CelebA() noise = RandomNoiseGenerator(latent_size, 'gaussian') pggan = PGGAN(G, D, data, noise, opts) pggan.train()
transform = transforms.Compose([ transforms.Resize(resize), transforms.CenterCrop(resize), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) imgs = [] for imgfile in args.files: img = transform(Image.open(imgfile).convert("RGB")) imgs.append(img) imgs = torch.stack(imgs, 0).to(device) g_ema = Generator(args.size, 512, 8) g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False) g_ema.eval() g_ema = g_ema.to(device) with torch.no_grad(): noise_sample = torch.randn(n_mean_latent, 512, device=device) latent_out = g_ema.style(noise_sample) latent_mean = latent_out.mean(0) latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent)**0.5 percept = lpips.PerceptualLoss(model="net-lin", net="vgg", use_gpu=device.startswith("cuda"))
def main(args): # by default we only consider single gpu inference assert (len(args.gpu) == 1) os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # load data data_loader_val, num_test = get_loader(args, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, training=False) print('finished data loading') # Generator colorguide = True if args.nocolor: colorguide = False netG = Generator(lambdas=None, colorguide=colorguide, input_nc=1, output_nc=1) netG.load_state_dict(torch.load(args.model_path)) if torch.cuda.is_available(): netG = netG.cuda() out_path = args.out_path checkpath(out_path) predictions_fid_real = [] predictions_fid_fake = [] fid_model = InceptionV3().cuda() fid_model.eval() Perceptual = PNet().cuda() avg_ssim = 0 lpips = 0 # validate on test set, TODO: test with single color guide image with torch.no_grad(): netG.eval() for i, (img_real, wf_real, color_real) in enumerate(data_loader_val, 0): img_real = img_real.cuda() wf_real = wf_real.cuda() if colorguide: color_real = color_real.cuda() # in case we are in the last interation batch_size = img_real.size(0) img_fake, wf_fake, _, _, _, _, _ = netG(trainG=False, img_real=None, wf_real=wf_real, color_real=color_real) ssim_score = ssim(img_real, img_fake).item() * batch_size avg_ssim += ssim_score lpips += Perceptual(img_real, img_fake) * batch_size # TODO: save generated wireframes save_singleimages(img_fake, out_path, i * args.batch_size, args.img_size) pred_fid_real = fid_model(img_real)[0] pred_fid_fake = fid_model(img_fake)[0] predictions_fid_real.append( pred_fid_real.data.cpu().numpy().reshape(batch_size, -1)) predictions_fid_fake.append( pred_fid_fake.data.cpu().numpy().reshape(batch_size, -1)) print('SSIM: {:6f}'.format(avg_ssim / num_test)) print('LPIPS: {:6f}'.format(lpips / num_test)) predictions_fid_real = np.concatenate(predictions_fid_real, 0) predictions_fid_fake = np.concatenate(predictions_fid_fake, 0) fid = compute_fid_score(predictions_fid_fake, predictions_fid_real) print('FID: {:6f}'.format(fid))
if args.gif and args.num_samples < 2: raise ValueError('for GIF num_samples must be greater than 1') if not args.out_path: out_path = 'ProGAN-Anime-Faces' else: out_path = args.out_path if args.device == 'cuda': if torch.cuda.is_available(): device = torch.device('cuda') else: device = torch.device('cpu') gen = Generator(cfg.Z_DIMENSION, cfg.IN_CHANNELS, cfg.CHANNELS_IMG).to(device) load_gen(gen, args.path_ckpt, device) gen.eval() alpha = 1 step = int(log2(cfg.START_TRAIN_IMG_SIZE / 4)) if args.grid: noise = get_random_noise(args.num_samples, args.z_size, device) print("==> Generate IMAGE GRID...") output = gen(noise, alpha, step) show_batch(output, out_path, num_samples=args.num_samples, figsize=(args.img_size, args.img_size)) elif args.gif: noise = get_random_noise(args.num_samples, args.z_size, device)
"--out_prefix", type=str, default="factor", help="filename prefix to result samples", ) parser.add_argument( "factor", type=str, help="name of the closed form factorization result factor file", ) args = parser.parse_args() eigvec = torch.load(args.factor)["eigvec"].to(args.device) ckpt = torch.load(args.ckpt) g = Generator(args.size, 512, 8, channel_multiplier=args.channel_multiplier).to(args.device) g.load_state_dict(ckpt["g_ema"], strict=False) trunc = g.mean_latent(4096) latent = torch.randn(args.n_sample, 512, device=args.device) latent = g.get_latent(latent) direction = args.degree * eigvec[:, args.index].unsqueeze(0) img, _ = g( [latent], truncation=args.truncation, truncation_latent=trunc, input_is_latent=True, )
parser.add_argument( "--inception", type=str, default=None, required=True, help="path to precomputed inception embedding", ) parser.add_argument( "ckpt", metavar="CHECKPOINT", help="path to generator checkpoint" ) args = parser.parse_args() ckpt = torch.load(args.ckpt) g = Generator(args.size, 512, 8).to(device) g.load_state_dict(ckpt["g_ema"]) g = nn.DataParallel(g) g.eval() if args.truncation < 1: with torch.no_grad(): mean_latent = g.mean_latent(args.truncation_mean) else: mean_latent = None inception = nn.DataParallel(load_patched_inception_v3()).to(device) inception.eval() features = extract_feature_from_samples(
args.distributed = n_gpu > 1 if args.distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") synchronize() args.latent = 512 args.n_mlp = 8 args.start_iter = 0 generator = Generator( args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier).to(device) discriminator = Discriminator( args.size, channel_multiplier=args.channel_multiplier).to(device) g_ema = Generator(args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier).to(device) g_ema.eval() accumulate(g_ema, generator, 0) g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1) d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1) g_optim = optim.Adam(
from data.dataloader import make_dataloader ImageFile.LOAD_TRUNCATED_IMAGES = True opt = TrainOptions().parse() localtime = time.asctime(time.localtime(time.time())) start_time = time.time() train_loader, valid_loader = make_dataloader(opt) # Decide which device we want to run on device = torch.device("cuda:0" if ( torch.cuda.is_available() and opt.ngpu > 0) else "cpu") # Create the generator netG = Generator().to(device) if (device.type == 'cuda') and (opt.ngpu > 1): netG = nn.DataParallel(netG, list(range(opt.ngpu))) netG.apply(weights_init) print(netG) # Create the Discriminator netD = Discriminator().to(device) if (device.type == 'cuda') and (opt.ngpu > 1): netD = nn.DataParallel(netD, list(range(opt.ngpu))) netD.apply(weights_init) print(netD) # Initialize BCELoss function criterion = nn.BCELoss() l1_loss = nn.L1Loss(reduction='sum')
opt = parser.parse_args() # dataset and dataloader path = "D:\\python\\worksapce\\P_github_project_INFOgan_mnist_ten\\data" dataset = MyDataset(path, opt) dataloader = DataLoader(dataset, batch_size=opt.batchSize, shuffle=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") ngpu = int(opt.ngpu) nz = int(opt.nz) ngf = int(opt.ngf) ndf = int(opt.ndf) nc = 1 netG = Generator(ngpu, nz, ngf, nc).to(device) netD = Discriminator(ngpu, nc, ndf).to(device) netQ = QNet(ngpu, nc, ndf).to(device) criterion = nn.BCELoss() qnetloss = QNetLoss() optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerQ = optim.Adam(netQ.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) if __name__ == "__main__": netG.apply(weights_init) netD.apply(weights_init) fixed_z = Variable(torch.randn(opt.batchSize, nz - 10, 1, 1)).to(device)
data_range = 10 batch_size = args.batch_size num_epochs = args.max_epoch input_dim = 1 hidden_dim = 32 output_dim = 1 num_epochs = 100000 num_epochs_pre = 500 learning_rate = 0.03 # Samples data = DataDistribution(mu, sigma) gen = NoiseDistribution(data_range) # Models G = Generator(input_dim, hidden_dim, output_dim) D = Discriminator(input_dim, hidden_dim, output_dim) # Loss function criterion = torch.nn.BCELoss() # optimizer optimizer = torch.optim.SGD(D.parameters(), lr=learning_rate) D_pre_losses = [] num_samples_pre = 5000 num_bins_pre = 100 for epoch in range(num_epochs_pre): # Generate samples d = data.sample(num_samples_pre) histc, edges = np.histogram(d, num_bins_pre, density=True)
# visdom setting vissum = VisdomSummary(port=FG.vis_port, env=FG.vis_env) # Dimensionality of the latent vector. latent_size = 512 # Use sigmoid activation for the last layer? sigmoid_at_end = FG.gan in ['lsgan', 'gan'] if hasattr(FG, 'no_tanh'): tanh_at_end = False else: tanh_at_end = True G = Generator(num_channels=3, latent_size=latent_size, resolution=FG.target_resol, fmap_max=latent_size, fmap_base=8192, tanh_at_end=tanh_at_end).to(device) D = Discriminator(num_channels=3, mbstat_avg=FG.mbstat_avg, resolution=FG.target_resol, fmap_max=latent_size, fmap_base=8192, sigmoid_at_end=sigmoid_at_end).to(device) print(G) print(D) # exit() if len(FG.gpu) != 1: G = torch.nn.DataParallel(G, FG.gpu)
def test(args): device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu") transform = [transforms.Resize((256, 256), Image.BICUBIC), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] test_data_loader = dataset( root_dir = args.DATA.data_path, dataset = args.DATA.dataset, mode = 'test', direction = args.DATA.direction, transform=transform) test_loader = DataLoader(test_data_loader, batch_size=1, shuffle=True, num_workers=4) "Define Models" generator = Generator(args.MODEL.ngf, args.MODEL.input_nc, args.MODEL.output_nc) "Load weights" checkpoint = args.EVALUATION.evaluation_path checkpoint = os.path.join(checkpoint, "best.pth.tar") load(generator, checkpoint) generator.to(device) L1_loss = torch.nn.L1Loss() t_loss = 0 generator.eval() for i, batch in enumerate(test_loader): input = batch['input'].to(device) target = batch['target'].to(device) "Discriminator Training" with torch.no_grad(): output = generator.forward(input) test_loss = L1_loss(output,target) t_loss += test_loss if args.EVALUATION.plot: "Display Images" input = input.detach().cpu().numpy() input = input.squeeze().transpose((1, 2, 0)) output = output.detach().cpu().numpy() output = output.squeeze().transpose((1, 2, 0)) fig = plt.figure() ax1 = fig.add_subplot(1, 2, 1) ax1.set_title('input') ax1.imshow(input) ax2 = fig.add_subplot(1, 2, 2) ax2.set_title('output') ax2.imshow(output) plt.pause(1) print("Iter:", i, "L1 loss:", test_loss.item()) print("Final Test Loss: ", t_loss.item()/len(test_loader))
device = "cpu" print(f"=> Called with args {args.__dict__}") print(f"=> Config params {cfg.__dict__}") print(f"=> Run on device {device}") # define dataset and dataloader dataset = AnimeFacesDataset(args.data_path) cfg.DATASET_SIZE = len(dataset) dataloader = DataLoader(dataset, batch_size=cfg.BATCH_SIZE, shuffle=True, num_workers=2, drop_last=True, pin_memory=True) # define models gen = Generator(cfg.Z_DIMENSION, cfg.CHANNELS_IMG, cfg.FEATURES_GEN).to(device) critic = Critic(cfg.CHANNELS_IMG, cfg.FEATURES_DISC).to(device) if args.checkpoint_path: opt_gen = optim.Adam(gen.parameters(), lr=cfg.LEARNING_RATE, betas=(0.5, 0.999)) opt_critic = optim.Adam(critic.parameters(), lr=cfg.LEARNING_RATE, betas=(0.5, 0.999)) cp = torch.load(args.checkpoint_path) start_epoch, end_epoch, fixed_noise = load_checkpoint( cp, gen, critic, opt_gen, opt_critic) cfg.NUM_EPOCHS = end_epoch else: print("=> Init default weights of models and fixed noise")
batch_size=params.batch_size, shuffle=True) # Test data test_data = DatasetFromFolder(data_dir, val_file, params.img_types, transform=transform) test_data_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=params.batch_size, shuffle=False) # test_input, test_target = test_data_loader.__iter__().__next__() # Models torch.cuda.set_device(params.gpu) G = Generator(3, params.ngf, 3) D = Discriminator(6, params.ndf, 1) G.cuda() D.cuda() G.normal_weight_init(mean=0.0, std=0.02) D.normal_weight_init(mean=0.0, std=0.02) slim_params, insnorm_params = [], [] for name, param in G.named_parameters(): if param.requires_grad and name.endswith( 'weight') and 'insnorm_conv' in name: insnorm_params.append(param) if len(slim_params) % 2 == 0: slim_params.append(param[:len(param) // 2]) else: slim_params.append(param[len(param) // 2:])
G_path = os.path.join('./checkpoints', opt.checkpoint, 'weight', 'netG_epoch_' + str(opt.startEpoch) + '.pth') D_path = os.path.join('./checkpoints', opt.checkpoint, 'weight', 'netD_epoch_' + str(opt.startEpoch) + '.pth') print(opt) driving_video_loader = make_test_dataloader(opt) # Decide which device we want to run on device = torch.device("cuda:0" if ( torch.cuda.is_available() and ngpu > 0) else "cpu") # Create the generator netG = Generator(ngf).to(device) if (device.type == 'cuda') and (ngpu > 1): netG = nn.DataParallel(netG, list(range(ngpu))) print(netG) # load model netG.load_state_dict(torch.load(G_path)) print('load %s !' % G_path) # source image source_image = Image.open(opt.sourceImage).convert('RGB') source_image_transform = transforms.Compose([ transforms.Resize((opt.imageSize, opt.imageSize)), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])
if args.fid: fid_model = InceptionV3FID(device) else: fid_model = None # defining dataset and dataloader dataset = ImgFolderDataset(args.data_path) dataloader = DataLoader(dataset, batch_size=cfg.BATCH_SIZE, shuffle=True, num_workers=2, drop_last=True, pin_memory=True) # defining models gen = Generator(img_size=cfg.IMG_SIZE, in_channels=cfg.IN_CHANNELS, img_channels=cfg.CHANNELS_IMG, z_dim=cfg.Z_DIMENSION).to(device) dis = Discriminator(img_size=cfg.IMG_SIZE, img_channels=cfg.CHANNELS_IMG).to(device) # defining optimizers opt_gen = optim.Adam(params=gen.parameters(), lr=cfg.LEARNING_RATE, betas=(0.0, 0.99)) opt_dis = optim.Adam(params=dis.parameters(), lr=cfg.LEARNING_RATE, betas=(0.0, 0.99)) # defining gradient scalers for automatic mixed precision scaler_gen = torch.cuda.amp.GradScaler() scaler_dis = torch.cuda.amp.GradScaler() if args.checkpoint: