Spaces:
Runtime error
Runtime error
| import argparse | |
| import math | |
| import random | |
| import os | |
| from util import * | |
| import numpy as np | |
| import torch | |
| torch.backends.cudnn.benchmark = True | |
| from torch import nn, autograd | |
| from torch import optim | |
| from torch.nn import functional as F | |
| from torch.utils import data | |
| import torch.distributed as dist | |
| from torchvision import transforms, utils | |
| from tqdm import tqdm | |
| from torch.optim import lr_scheduler | |
| import copy | |
| import kornia.augmentation as K | |
| import kornia | |
| import lpips | |
| from model import * | |
| from dataset import ImageFolder | |
| from distributed import ( | |
| get_rank, | |
| synchronize, | |
| reduce_loss_dict, | |
| reduce_sum, | |
| get_world_size, | |
| ) | |
| mse_criterion = nn.MSELoss() | |
| def test(args, genA2B, genB2A, testA_loader, testB_loader, name, step): | |
| testA_loader = iter(testA_loader) | |
| testB_loader = iter(testB_loader) | |
| with torch.no_grad(): | |
| test_sample_num = 16 | |
| genA2B.eval(), genB2A.eval() | |
| A2B = [] | |
| B2A = [] | |
| for i in range(test_sample_num): | |
| real_A = testA_loader.next() | |
| real_B = testB_loader.next() | |
| real_A, real_B = real_A.cuda(), real_B.cuda() | |
| A2B_content, A2B_style = genA2B.encode(real_A) | |
| B2A_content, B2A_style = genB2A.encode(real_B) | |
| if i % 2 == 0: | |
| A2B_mod1 = torch.randn([1, args.latent_dim]).cuda() | |
| B2A_mod1 = torch.randn([1, args.latent_dim]).cuda() | |
| A2B_mod2 = torch.randn([1, args.latent_dim]).cuda() | |
| B2A_mod2 = torch.randn([1, args.latent_dim]).cuda() | |
| fake_B2B, _, _ = genA2B(real_B) | |
| fake_A2A, _, _ = genB2A(real_A) | |
| colsA = [real_A, fake_A2A] | |
| colsB = [real_B, fake_B2B] | |
| fake_A2B_1 = genA2B.decode(A2B_content, A2B_mod1) | |
| fake_B2A_1 = genB2A.decode(B2A_content, B2A_mod1) | |
| fake_A2B_2 = genA2B.decode(A2B_content, A2B_mod2) | |
| fake_B2A_2 = genB2A.decode(B2A_content, B2A_mod2) | |
| fake_A2B_3 = genA2B.decode(A2B_content, B2A_style) | |
| fake_B2A_3 = genB2A.decode(B2A_content, A2B_style) | |
| colsA += [fake_A2B_3, fake_A2B_1, fake_A2B_2] | |
| colsB += [fake_B2A_3, fake_B2A_1, fake_B2A_2] | |
| fake_A2B2A, _, _ = genB2A(fake_A2B_3, A2B_style) | |
| fake_B2A2B, _, _ = genA2B(fake_B2A_3, B2A_style) | |
| colsA.append(fake_A2B2A) | |
| colsB.append(fake_B2A2B) | |
| fake_A2B2A, _, _ = genB2A(fake_A2B_1, A2B_style) | |
| fake_B2A2B, _, _ = genA2B(fake_B2A_1, B2A_style) | |
| colsA.append(fake_A2B2A) | |
| colsB.append(fake_B2A2B) | |
| fake_A2B2A, _, _ = genB2A(fake_A2B_2, A2B_style) | |
| fake_B2A2B, _, _ = genA2B(fake_B2A_2, B2A_style) | |
| colsA.append(fake_A2B2A) | |
| colsB.append(fake_B2A2B) | |
| fake_A2B2A, _, _ = genB2A(fake_A2B_1) | |
| fake_B2A2B, _, _ = genA2B(fake_B2A_1) | |
| colsA.append(fake_A2B2A) | |
| colsB.append(fake_B2A2B) | |
| colsA = torch.cat(colsA, 2).detach().cpu() | |
| colsB = torch.cat(colsB, 2).detach().cpu() | |
| A2B.append(colsA) | |
| B2A.append(colsB) | |
| A2B = torch.cat(A2B, 0) | |
| B2A = torch.cat(B2A, 0) | |
| utils.save_image(A2B, f'{im_path}/{name}_A2B_{str(step).zfill(6)}.jpg', normalize=True, range=(-1, 1), nrow=16) | |
| utils.save_image(B2A, f'{im_path}/{name}_B2A_{str(step).zfill(6)}.jpg', normalize=True, range=(-1, 1), nrow=16) | |
| genA2B.train(), genB2A.train() | |
| def train(args, trainA_loader, trainB_loader, testA_loader, testB_loader, G_A2B, G_B2A, D_A, D_B, G_optim, D_optim, device): | |
| G_A2B.train(), G_B2A.train(), D_A.train(), D_B.train() | |
| trainA_loader = sample_data(trainA_loader) | |
| trainB_loader = sample_data(trainB_loader) | |
| G_scheduler = lr_scheduler.StepLR(G_optim, step_size=100000, gamma=0.5) | |
| D_scheduler = lr_scheduler.StepLR(D_optim, step_size=100000, gamma=0.5) | |
| pbar = range(args.iter) | |
| if get_rank() == 0: | |
| pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0.1) | |
| loss_dict = {} | |
| mean_path_length_A2B = 0 | |
| mean_path_length_B2A = 0 | |
| if args.distributed: | |
| G_A2B_module = G_A2B.module | |
| G_B2A_module = G_B2A.module | |
| D_A_module = D_A.module | |
| D_B_module = D_B.module | |
| D_L_module = D_L.module | |
| else: | |
| G_A2B_module = G_A2B | |
| G_B2A_module = G_B2A | |
| D_A_module = D_A | |
| D_B_module = D_B | |
| D_L_module = D_L | |
| for idx in pbar: | |
| i = idx + args.start_iter | |
| if i > args.iter: | |
| print('Done!') | |
| break | |
| ori_A = next(trainA_loader) | |
| ori_B = next(trainB_loader) | |
| if isinstance(ori_A, list): | |
| ori_A = ori_A[0] | |
| if isinstance(ori_B, list): | |
| ori_B = ori_B[0] | |
| ori_A = ori_A.to(device) | |
| ori_B = ori_B.to(device) | |
| aug_A = aug(ori_A) | |
| aug_B = aug(ori_B) | |
| A = aug(ori_A[[np.random.randint(args.batch)]].expand_as(ori_A)) | |
| B = aug(ori_B[[np.random.randint(args.batch)]].expand_as(ori_B)) | |
| if i % args.d_reg_every == 0: | |
| aug_A.requires_grad = True | |
| aug_B.requires_grad = True | |
| A2B_content, A2B_style = G_A2B.encode(A) | |
| B2A_content, B2A_style = G_B2A.encode(B) | |
| # get new style | |
| aug_A2B_style = G_B2A.style_encode(aug_B) | |
| aug_B2A_style = G_A2B.style_encode(aug_A) | |
| rand_A2B_style = torch.randn([args.batch, args.latent_dim]).to(device).requires_grad_() | |
| rand_B2A_style = torch.randn([args.batch, args.latent_dim]).to(device).requires_grad_() | |
| # styles | |
| idx = torch.randperm(2*args.batch) | |
| input_A2B_style = torch.cat([rand_A2B_style, aug_A2B_style], 0)[idx][:args.batch] | |
| idx = torch.randperm(2*args.batch) | |
| input_B2A_style = torch.cat([rand_B2A_style, aug_B2A_style], 0)[idx][:args.batch] | |
| fake_A2B = G_A2B.decode(A2B_content, input_A2B_style) | |
| fake_B2A = G_B2A.decode(B2A_content, input_B2A_style) | |
| # train disc | |
| real_A_logit = D_A(aug_A) | |
| real_B_logit = D_B(aug_B) | |
| real_L_logit1 = D_L(rand_A2B_style) | |
| real_L_logit2 = D_L(rand_B2A_style) | |
| fake_B_logit = D_B(fake_A2B.detach()) | |
| fake_A_logit = D_A(fake_B2A.detach()) | |
| fake_L_logit1 = D_L(aug_A2B_style.detach()) | |
| fake_L_logit2 = D_L(aug_B2A_style.detach()) | |
| # global loss | |
| D_loss = d_logistic_loss(real_A_logit, fake_A_logit) +\ | |
| d_logistic_loss(real_B_logit, fake_B_logit) +\ | |
| d_logistic_loss(real_L_logit1, fake_L_logit1) +\ | |
| d_logistic_loss(real_L_logit2, fake_L_logit2) | |
| loss_dict['D_adv'] = D_loss | |
| if i % args.d_reg_every == 0: | |
| r1_A_loss = d_r1_loss(real_A_logit, aug_A) | |
| r1_B_loss = d_r1_loss(real_B_logit, aug_B) | |
| r1_L_loss = d_r1_loss(real_L_logit1, rand_A2B_style) + d_r1_loss(real_L_logit2, rand_B2A_style) | |
| r1_loss = r1_A_loss + r1_B_loss + r1_L_loss | |
| D_r1_loss = (args.r1 / 2 * r1_loss * args.d_reg_every) | |
| D_loss += D_r1_loss | |
| D_optim.zero_grad() | |
| D_loss.backward() | |
| D_optim.step() | |
| #Generator | |
| # adv loss | |
| fake_B_logit = D_B(fake_A2B) | |
| fake_A_logit = D_A(fake_B2A) | |
| fake_L_logit1 = D_L(aug_A2B_style) | |
| fake_L_logit2 = D_L(aug_B2A_style) | |
| lambda_adv = (1, 1, 1) | |
| G_adv_loss = 1 * (g_nonsaturating_loss(fake_A_logit, lambda_adv) +\ | |
| g_nonsaturating_loss(fake_B_logit, lambda_adv) +\ | |
| 2*g_nonsaturating_loss(fake_L_logit1, (1,)) +\ | |
| 2*g_nonsaturating_loss(fake_L_logit2, (1,))) | |
| # style consis loss | |
| G_con_loss = 50 * (A2B_style.var(0, unbiased=False).sum() + B2A_style.var(0, unbiased=False).sum()) | |
| # cycle recon | |
| A2B2A_content, A2B2A_style = G_B2A.encode(fake_A2B) | |
| B2A2B_content, B2A2B_style = G_A2B.encode(fake_B2A) | |
| fake_A2B2A = G_B2A.decode(A2B2A_content, shuffle_batch(A2B_style)) | |
| fake_B2A2B = G_A2B.decode(B2A2B_content, shuffle_batch(B2A_style)) | |
| G_cycle_loss = 20 * (F.mse_loss(fake_A2B2A, A) + F.mse_loss(fake_B2A2B, B)) | |
| lpips_loss = 10 * (lpips_fn(fake_A2B2A, A).mean() + lpips_fn(fake_B2A2B, B).mean()) #10 for anime | |
| # style reconstruction | |
| G_style_loss = 5 * (mse_criterion(A2B2A_style, input_A2B_style) +\ | |
| mse_criterion(B2A2B_style, input_B2A_style)) | |
| G_loss = G_adv_loss + G_cycle_loss + G_con_loss + lpips_loss + G_style_loss | |
| loss_dict['G_adv'] = G_adv_loss | |
| loss_dict['G_con'] = G_con_loss | |
| loss_dict['G_cycle'] = G_cycle_loss | |
| loss_dict['lpips'] = lpips_loss | |
| G_optim.zero_grad() | |
| G_loss.backward() | |
| G_optim.step() | |
| G_scheduler.step() | |
| D_scheduler.step() | |
| accumulate(G_A2B_ema, G_A2B_module) | |
| accumulate(G_B2A_ema, G_B2A_module) | |
| loss_reduced = reduce_loss_dict(loss_dict) | |
| D_adv_loss_val = loss_reduced['D_adv'].mean().item() | |
| G_adv_loss_val = loss_reduced['G_adv'].mean().item() | |
| G_cycle_loss_val = loss_reduced['G_cycle'].mean().item() | |
| G_con_loss_val = loss_reduced['G_con'].mean().item() | |
| lpips_val = loss_reduced['lpips'].mean().item() | |
| if get_rank() == 0: | |
| pbar.set_description( | |
| ( | |
| f'Dadv: {D_adv_loss_val:.2f}; lpips: {lpips_val:.2f} ' | |
| f'Gadv: {G_adv_loss_val:.2f}; Gcycle: {G_cycle_loss_val:.2f}; GMS: {G_con_loss_val:.2f} {G_style_loss.item():.2f}' | |
| ) | |
| ) | |
| if i % 1000 == 0: | |
| with torch.no_grad(): | |
| test(args, G_A2B, G_B2A, testA_loader, testB_loader, 'normal', i) | |
| test(args, G_A2B_ema, G_B2A_ema, testA_loader, testB_loader, 'ema', i) | |
| if (i+1) % 2000 == 0: | |
| torch.save( | |
| { | |
| 'G_A2B': G_A2B_module.state_dict(), | |
| 'G_B2A': G_B2A_module.state_dict(), | |
| 'G_A2B_ema': G_A2B_ema.state_dict(), | |
| 'G_B2A_ema': G_B2A_ema.state_dict(), | |
| 'D_A': D_A_module.state_dict(), | |
| 'D_B': D_B_module.state_dict(), | |
| 'D_L': D_L_module.state_dict(), | |
| 'G_optim': G_optim.state_dict(), | |
| 'D_optim': D_optim.state_dict(), | |
| 'iter': i, | |
| }, | |
| os.path.join(model_path, 'ck.pt'), | |
| ) | |
| if __name__ == '__main__': | |
| device = 'cuda' | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--iter', type=int, default=300000) | |
| parser.add_argument('--batch', type=int, default=4) | |
| parser.add_argument('--n_sample', type=int, default=64) | |
| parser.add_argument('--size', type=int, default=256) | |
| parser.add_argument('--r1', type=float, default=10) | |
| parser.add_argument('--lambda_cycle', type=int, default=1) | |
| parser.add_argument('--path_regularize', type=float, default=2) | |
| parser.add_argument('--path_batch_shrink', type=int, default=2) | |
| parser.add_argument('--d_reg_every', type=int, default=16) | |
| parser.add_argument('--g_reg_every', type=int, default=4) | |
| parser.add_argument('--mixing', type=float, default=0.9) | |
| parser.add_argument('--ckpt', type=str, default=None) | |
| parser.add_argument('--lr', type=float, default=2e-3) | |
| parser.add_argument('--local_rank', type=int, default=0) | |
| parser.add_argument('--num_down', type=int, default=3) | |
| parser.add_argument('--name', type=str, required=True) | |
| parser.add_argument('--d_path', type=str, required=True) | |
| parser.add_argument('--latent_dim', type=int, default=8) | |
| parser.add_argument('--lr_mlp', type=float, default=0.01) | |
| parser.add_argument('--n_res', type=int, default=1) | |
| args = parser.parse_args() | |
| n_gpu = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1 | |
| args.distributed = False | |
| if args.distributed: | |
| torch.cuda.set_device(args.local_rank) | |
| torch.distributed.init_process_group(backend='nccl', init_method='env://') | |
| synchronize() | |
| save_path = f'./{args.name}' | |
| im_path = os.path.join(save_path, 'sample') | |
| model_path = os.path.join(save_path, 'checkpoint') | |
| os.makedirs(im_path, exist_ok=True) | |
| os.makedirs(model_path, exist_ok=True) | |
| args.n_mlp = 5 | |
| args.start_iter = 0 | |
| G_A2B = Generator( args.size, args.num_down, args.latent_dim, args.n_mlp, lr_mlp=args.lr_mlp, n_res=args.n_res).to(device) | |
| D_A = Discriminator(args.size).to(device) | |
| G_B2A = Generator( args.size, args.num_down, args.latent_dim, args.n_mlp, lr_mlp=args.lr_mlp, n_res=args.n_res).to(device) | |
| D_B = Discriminator(args.size).to(device) | |
| D_L = LatDiscriminator(args.latent_dim).to(device) | |
| lpips_fn = lpips.LPIPS(net='vgg').to(device) | |
| G_A2B_ema = copy.deepcopy(G_A2B).to(device).eval() | |
| G_B2A_ema = copy.deepcopy(G_B2A).to(device).eval() | |
| 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( list(G_A2B.parameters()) + list(G_B2A.parameters()), lr=args.lr, betas=(0, 0.99)) | |
| D_optim = optim.Adam( | |
| list(D_L.parameters()) + list(D_A.parameters()) + list(D_B.parameters()), | |
| lr=args.lr, betas=(0**d_reg_ratio, 0.99**d_reg_ratio)) | |
| if args.ckpt is not None: | |
| ckpt = torch.load(args.ckpt, map_location=lambda storage, loc: storage) | |
| try: | |
| ckpt_name = os.path.basename(args.ckpt) | |
| args.start_iter = int(os.path.splitext(ckpt_name)[0]) | |
| except ValueError: | |
| pass | |
| G_A2B.load_state_dict(ckpt['G_A2B']) | |
| G_B2A.load_state_dict(ckpt['G_B2A']) | |
| G_A2B_ema.load_state_dict(ckpt['G_A2B_ema']) | |
| G_B2A_ema.load_state_dict(ckpt['G_B2A_ema']) | |
| D_A.load_state_dict(ckpt['D_A']) | |
| D_B.load_state_dict(ckpt['D_B']) | |
| D_L.load_state_dict(ckpt['D_L']) | |
| G_optim.load_state_dict(ckpt['G_optim']) | |
| D_optim.load_state_dict(ckpt['D_optim']) | |
| args.start_iter = ckpt['iter'] | |
| if args.distributed: | |
| G_A2B = nn.parallel.DistributedDataParallel( | |
| G_A2B, | |
| device_ids=[args.local_rank], | |
| output_device=args.local_rank, | |
| broadcast_buffers=False, | |
| ) | |
| D_A = nn.parallel.DistributedDataParallel( | |
| D_A, | |
| device_ids=[args.local_rank], | |
| output_device=args.local_rank, | |
| broadcast_buffers=False, | |
| ) | |
| G_B2A = nn.parallel.DistributedDataParallel( | |
| G_B2A, | |
| device_ids=[args.local_rank], | |
| output_device=args.local_rank, | |
| broadcast_buffers=False, | |
| ) | |
| D_B = nn.parallel.DistributedDataParallel( | |
| D_B, | |
| device_ids=[args.local_rank], | |
| output_device=args.local_rank, | |
| broadcast_buffers=False, | |
| ) | |
| D_L = nn.parallel.DistributedDataParallel( | |
| D_L, | |
| device_ids=[args.local_rank], | |
| output_device=args.local_rank, | |
| broadcast_buffers=False, | |
| ) | |
| train_transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), inplace=True) | |
| ]) | |
| test_transform = transforms.Compose([ | |
| transforms.Resize((args.size, args.size)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), inplace=True) | |
| ]) | |
| aug = nn.Sequential( | |
| K.RandomAffine(degrees=(-20,20), scale=(0.8, 1.2), translate=(0.1, 0.1), shear=0.15), | |
| kornia.geometry.transform.Resize(256+30), | |
| K.RandomCrop((256,256)), | |
| K.RandomHorizontalFlip(), | |
| ) | |
| d_path = args.d_path | |
| trainA = ImageFolder(os.path.join(d_path, 'trainA'), train_transform) | |
| trainB = ImageFolder(os.path.join(d_path, 'trainB'), train_transform) | |
| testA = ImageFolder(os.path.join(d_path, 'testA'), test_transform) | |
| testB = ImageFolder(os.path.join(d_path, 'testB'), test_transform) | |
| trainA_loader = data.DataLoader(trainA, batch_size=args.batch, | |
| sampler=data_sampler(trainA, shuffle=True, distributed=args.distributed), drop_last=True, pin_memory=True, num_workers=5) | |
| trainB_loader = data.DataLoader(trainB, batch_size=args.batch, | |
| sampler=data_sampler(trainB, shuffle=True, distributed=args.distributed), drop_last=True, pin_memory=True, num_workers=5) | |
| testA_loader = data.DataLoader(testA, batch_size=1, shuffle=False) | |
| testB_loader = data.DataLoader(testB, batch_size=1, shuffle=False) | |
| train(args, trainA_loader, trainB_loader, testA_loader, testB_loader, G_A2B, G_B2A, D_A, D_B, G_optim, D_optim, device) | |