train.py

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)