if step % args.D_iter == 0: optim_G.zero_grad() z = torch.randn(args.batch_size, args.z_dim).to(device) loss_G = -net_D(net_G(z)).mean() loss_G.backward() optim_G.step() if step == 0: grid = (make_grid(real) + 1) / 2 train_writer.add_image('real sample', grid) if step == 0 or (step + 1) % args.sample_iter == 0: fake = net_G(sample_z).cpu() grid = (make_grid(fake) + 1) / 2 valid_writer.add_image('sample', grid, step) save_image(grid, os.path.join(log_dir, 'sample', '%d.png' % step)) if step == 0 or (step + 1) % 10000 == 0: torch.save(net_G.state_dict(), os.path.join(log_dir, 'G_%d.pt' % step)) score, _ = inception_score(valid_dataset, batch_size=64, cuda=True) valid_writer.add_scalar('Inception Score', score, step) score = fid_score(IgnoreLabelDataset(cifar10), valid_dataset, batch_size=64, cuda=True, normalize=True, r_cache='./.fid_cache/cifar10') valid_writer.add_scalar('FID Score', score, step)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--iterations', type=int, default=100000) parser.add_argument('--batch-size', type=int, default=64) parser.add_argument('--lr', type=float, default=1e-4) parser.add_argument('--consistency', type=float, default=10) parser.add_argument('--warm_up', type=float, default=1000) parser.add_argument('--gp_center', type=float, default=1) parser.add_argument('--gamma', type=int, default=10) parser.add_argument('--name', type=str, default='DCGAN_v2_CR_GP') parser.add_argument('--log-dir', type=str, default='log') parser.add_argument('--z-dim', type=int, default=128) parser.add_argument('--iter-G', type=int, default=3) parser.add_argument('--sample-iter', type=int, default=1000) parser.add_argument('--sample-size', type=int, default=64) args = parser.parse_args() log_dir = os.path.join(args.log_dir, args.name) device = torch.device('cuda') cifar10 = datasets.CIFAR10('./data', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataloader = torch.utils.data.DataLoader(cifar10, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True) net_G = Generator().to(device) net_D = Discriminator().to(device) optim_G = optim.Adam(net_G.parameters(), lr=args.lr, betas=(0.5, 0.999)) optim_D = optim.Adam(net_D.parameters(), lr=args.lr, betas=(0.5, 0.999)) train_writer = SummaryWriter(os.path.join(log_dir, 'train')) valid_writer = SummaryWriter(os.path.join(log_dir, 'valid')) real_label = torch.full((args.batch_size, 1), 1).to(device) fake_label = torch.full((args.batch_size, 1), 0).to(device) label = torch.cat([real_label, fake_label], dim=0) criteria = nn.BCEWithLogitsLoss() os.makedirs(os.path.join(log_dir, 'sample'), exist_ok=True) sample_z = torch.randn(args.sample_size, args.z_dim).to(device) valid_dataset = GenerativeDataset(net_G, args.z_dim, 10000, device) looper = loop(dataloader) consistency_transforms = transforms.Compose([ transforms.ToPILImage(mode='RGB'), transforms.RandomAffine(0, translate=(0.1, 0.1)), transforms.RandomHorizontalFlip(p=1.0), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) def consistency_transform_func(images): images = deepcopy(images) for idx, img in enumerate(images): images[idx] = consistency_transforms(img) return images cs_lambda = args.consistency with trange(args.iterations, dynamic_ncols=True) as pbar: for step in pbar: real, _ = next(looper) augment_real = consistency_transform_func(real) real = real.to(device) augment_real = augment_real.to(device) # update discriminator z = torch.randn(args.batch_size, args.z_dim).to(device) with torch.no_grad(): fake = net_G(z).detach() loss_gp = calc_gradient_penalty(net_D, real, fake, args.gp_center) loss_cs = ((net_D(real) - net_D(augment_real))**2).sum() pred_D = torch.cat([net_D(real), net_D(fake)]) loss_D = criteria( pred_D, label) + (args.gamma * loss_gp) + (cs_lambda * loss_cs) train_writer.add_scalar('regularization/weight', cs_lambda, step + 1) optim_D.zero_grad() loss_D.backward() optim_D.step() train_writer.add_scalar('loss', loss_D.item(), step + 1) if step % args.iter_G == 0: # update generator z = torch.randn(args.batch_size, args.z_dim).to(device) pred_G = net_D(net_G(z)) loss_G = criteria(pred_G, real_label) optim_G.zero_grad() loss_G.backward() optim_G.step() train_writer.add_scalar('loss/G', loss_G.item(), step + 1) pbar.set_postfix(loss_D='%.4f' % loss_D.item(), loss_G='%.4f' % loss_G.item(), loss_gp='%.4f' % loss_gp.item()) if step == 0: grid = (make_grid(real[:args.sample_size]) + 1) / 2 train_writer.add_image('real sample', grid) if step == 0 or (step + 1) % args.sample_iter == 0: fake = net_G(sample_z).cpu() grid = (make_grid(fake) + 1) / 2 valid_writer.add_image('sample', grid, step) save_image(grid, os.path.join(log_dir, 'sample', '%d.png' % step)) if step == 0 or (step + 1) % 10000 == 0: torch.save(net_G.state_dict(), os.path.join(log_dir, 'G_%d.pt' % step)) score, _ = inception_score(valid_dataset, batch_size=64, cuda=True) valid_writer.add_scalar('Inception Score', score, step) score = fid_score(IgnoreLabelDataset(cifar10), valid_dataset, batch_size=64, cuda=True, normalize=True, r_cache='./.fid_cache/cifar10') valid_writer.add_scalar('FID Score', score, step)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--iterations', type=int, default=200000) parser.add_argument('--batch-size', type=int, default=64) parser.add_argument('--lr', type=float, default=1e-4) parser.add_argument('--clip', type=float, default=1e-2) parser.add_argument('--name', type=str, default='WGAN_v2') parser.add_argument('--log-dir', type=str, default='log') parser.add_argument('--z-dim', type=int, default=128) parser.add_argument('--iter-D', type=int, default=5) parser.add_argument('--sample-iter', type=int, default=1000) parser.add_argument('--sample-size', type=int, default=64) args = parser.parse_args() log_dir = os.path.join(args.log_dir, args.name) device = torch.device('cuda') cifar10 = datasets.CIFAR10('./data', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataloader = torch.utils.data.DataLoader(cifar10, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True) net_G = Generator(args.z_dim).to(device) net_D = Discriminator().to(device) optim_G = optim.RMSprop(net_G.parameters(), lr=args.lr) optim_D = optim.RMSprop(net_D.parameters(), lr=args.lr) train_writer = SummaryWriter(os.path.join(log_dir, 'train')) valid_writer = SummaryWriter(os.path.join(log_dir, 'valid')) os.makedirs(os.path.join(log_dir, 'sample'), exist_ok=True) sample_z = torch.randn(args.sample_size, args.z_dim).to(device) valid_dataset = GenerativeDataset(net_G, args.z_dim, 10000, device) looper = loop(dataloader) with trange(args.iterations, dynamic_ncols=True) as pbar: for step in pbar: real, _ = next(looper) real = real.to(device) z = torch.randn(args.batch_size, args.z_dim).to(device) with torch.no_grad(): fake = net_G(z).detach() loss = -net_D(real).mean() + net_D(fake).mean() optim_D.zero_grad() loss.backward() optim_D.step() train_writer.add_scalar('loss', -loss.item(), step) pbar.set_postfix(loss='%.4f' % -loss.item()) for param in net_D.parameters(): param.data.clamp_(-args.clip, args.clip) if step % args.iter_D == 0: z = torch.randn(args.batch_size, args.z_dim).to(device) loss_G = -net_D(net_G(z)).mean() optim_G.zero_grad() loss_G.backward() optim_G.step() if step == 0: grid = (make_grid(real) + 1) / 2 train_writer.add_image('real sample', grid) if step == 0 or (step + 1) % args.sample_iter == 0: fake = net_G(sample_z).cpu() grid = (make_grid(fake) + 1) / 2 valid_writer.add_image('sample', grid, step) save_image(grid, os.path.join(log_dir, 'sample', '%d.png' % step)) if step == 0 or (step + 1) % 10000 == 0: torch.save(net_G.state_dict(), os.path.join(log_dir, 'G_%d.pt' % step)) score, _ = inception_score(valid_dataset, batch_size=64, cuda=True) valid_writer.add_scalar('Inception Score', score, step) score = fid_score(IgnoreLabelDataset(cifar10), valid_dataset, batch_size=64, cuda=True, normalize=True, r_cache='./.fid_cache/cifar10') valid_writer.add_scalar('FID Score', score, step)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--iterations', type=int, default=100000) parser.add_argument('--batch-size', type=int, default=64) parser.add_argument('--lr', type=float, default=2e-4) parser.add_argument('--consistency', type=float, default=10) parser.add_argument('--name', type=str, default='SNGAN_CR') parser.add_argument('--log-dir', type=str, default='log') parser.add_argument('--z-dim', type=int, default=128) parser.add_argument('--D-iter', type=int, default=5) parser.add_argument('--sample-iter', type=int, default=1000) parser.add_argument('--sample-size', type=int, default=64) args = parser.parse_args() log_dir = os.path.join(args.log_dir, args.name) device = torch.device('cuda') cifar10 = datasets.CIFAR10('./data', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataloader = torch.utils.data.DataLoader(cifar10, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True) net_G = Generator(args.z_dim).to(device) net_D = Discriminator().to(device) optim_G = optim.Adam(net_G.parameters(), lr=args.lr, betas=(0.5, 0.999)) optim_D = optim.Adam(net_D.parameters(), lr=args.lr, betas=(0.5, 0.999)) train_writer = SummaryWriter(os.path.join(log_dir, 'train')) valid_writer = SummaryWriter(os.path.join(log_dir, 'valid')) os.makedirs(os.path.join(log_dir, 'sample'), exist_ok=True) sample_z = torch.randn(args.sample_size, args.z_dim).to(device) valid_dataset = GenerativeDataset(net_G, args.z_dim, 10000, device) looper = loop(dataloader) consistency_transforms = transforms.Compose([ transforms.ToPILImage(mode='RGB'), transforms.RandomAffine(0, translate=(0.1, 0.1)), transforms.RandomHorizontalFlip(p=0.5), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) def consistency_transform_func(images): images = deepcopy(images) for idx, img in enumerate(images): images[idx] = consistency_transforms(img) return images cs_lambda = args.consistency with trange(args.iterations, dynamic_ncols=True) as pbar: for step in pbar: real, _ = next(looper) augment_real = consistency_transform_func(real) real = real.to(device) augment_real = augment_real.to(device) z = torch.randn(args.batch_size, args.z_dim).to(device) with torch.no_grad(): fake = net_G(z).detach() loss_real = torch.nn.functional.relu(1 - net_D(real)).mean() loss_fake = torch.nn.functional.relu(1 + net_D(fake)).mean() loss_cs = ((net_D(real) - net_D(augment_real))**2).sum() loss_D = loss_real + loss_fake + cs_lambda * loss_cs optim_D.zero_grad() loss_D.backward() optim_D.step() train_writer.add_scalar('loss', loss_D.item(), step) pbar.set_postfix(loss='%.4f' % loss_D.item()) if step % args.D_iter == 0: optim_G.zero_grad() z = torch.randn(args.batch_size, args.z_dim).to(device) loss_G = -net_D(net_G(z)).mean() loss_G.backward() optim_G.step() if step == 0: grid = (make_grid(real) + 1) / 2 train_writer.add_image('real sample', grid) if step == 0 or (step + 1) % args.sample_iter == 0: fake = net_G(sample_z).cpu() grid = (make_grid(fake) + 1) / 2 valid_writer.add_image('sample', grid, step) save_image(grid, os.path.join(log_dir, 'sample', '%d.png' % step)) if step == 0 or (step + 1) % 10000 == 0: torch.save(net_G.state_dict(), os.path.join(log_dir, 'G_%d.pt' % step)) score, _ = inception_score(valid_dataset, batch_size=64, cuda=True) valid_writer.add_scalar('Inception Score', score, step) score = fid_score(IgnoreLabelDataset(cifar10), valid_dataset, batch_size=64, cuda=True, normalize=True, r_cache='./.fid_cache/cifar10') valid_writer.add_scalar('FID Score', score, step)