train_rcvitgan_Ds.py

from argparse import ArgumentParser
from pathlib import Path
import shutil

import imageio
def silence_imageio_warning(*args, **kwargs):
    pass
imageio.core.util._precision_warn = silence_imageio_warning

import gin
import numpy as np
import torch
import torch.nn as nn
from torch import autograd
import torch.optim as optim
from torch.utils.data import DataLoader

#from evaluate.gan import FIDScore, FixedSampleGeneration, ImageGrid
from datasets import get_dataset
from augment import get_augment
from models.gan import get_architecture
from utils import cycle,cycle3,cycle4

from training.gan import setup
from utils import Logger
from utils import count_parameters
from utils import accumulate
from utils import set_grad
#from data_parallel1 import BalancedDataParallel

# import for gin binding
import penalty

import wandb
import time

# import for evaluation
#from evaluate.gan import FIDScore, FixedSampleGeneration, ImageGrid

#from torchvision import datasets, transforms
from mydiscriminator import ResidualDiscriminatorP#,Pix2PixDiscriminator
from fid_score import my_fid_score

from tensorboardX import SummaryWriter
writer=SummaryWriter('out/log_dino_noSLN_smoothL1_Ds')

from collections import OrderedDict
from ignite.engine import *
from ignite.handlers import *
from ignite.metrics import *
from ignite.utils import *
from ignite.contrib.metrics.regression import *
from ignite.contrib.metrics import *
try:
    from third_party.fid.inception import InceptionV3
except ImportError:
    from inception import InceptionV3
def eval_step(engine, batch):
    return batch
default_evaluator = Engine(eval_step)
default_model = nn.Sequential(OrderedDict([
    ('base', nn.Linear(4, 2)),
    ('fc', nn.Linear(2, 1))
]))
# # wrapper class as feature_extractor
# class WrapperInceptionV3(nn.Module):

#     def __init__(self, fid_incv3):
#         super().__init__()
#         self.fid_incv3 = fid_incv3

#     @torch.no_grad()
#     def forward(self, x):
#         y = self.fid_incv3(x)
#         y = y[0]
#         y = y[:, :, 0, 0]
#         return y
    
# block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[192]
# model = InceptionV3([block_idx]).cuda()
# # wrapper model to pytorch_fid model
# wrapper_model = WrapperInceptionV3(model)
# wrapper_model.eval()

from color_harmonization import ch_loss
import torch.nn.functional as F
import warnings
warnings.filterwarnings("ignore",category=DeprecationWarning)
def parse_args():
    """Training script for StyleGAN2."""

    parser = ArgumentParser(description='Training script: StyleGAN2 with DataParallel.')
    parser.add_argument('gin_config', type=str, help='Path to the gin configuration file')
    parser.add_argument('architecture', type=str, help='Architecture')

    parser.add_argument('--mode', default='std', type=str, help='Training mode (default: std)')
    parser.add_argument('--penalty', default='none', type=str, help='Penalty (default: none)')
    parser.add_argument('--aug', default='none', type=str, help='Augmentation (default: hfrt)')
    parser.add_argument('--use_warmup', action='store_true', help='Use warmup strategy on LR')
    parser.add_argument('--workers', default=8, type=int, metavar='N',
                        help='number of data loading workers (default: 0)')

    parser.add_argument('--temp', default=0.1, type=float,
                        help='Temperature hyperparameter for contrastive losses')
    parser.add_argument('--lbd_a', default=1.0, type=float,
                        help='Relative strength of the fake loss of ContraD')

    # Options for StyleGAN2 training
    parser.add_argument('--no_lazy', action='store_true',
                        help='Do not use lazy regularization')
    parser.add_argument("--d_reg_every", type=int, default=16,
                        help='Interval of applying R1 when lazy regularization is used')
    parser.add_argument("--lbd_r1", type=float, default=10, help='R1 regularization')
    parser.add_argument('--style_mix', default=0.9, type=float, help='Style mixing regularization')
    parser.add_argument('--halflife_k', default=20, type=int,
                        help='Half-life of exponential moving average in thousands of images')
    parser.add_argument('--ema_start_k', default=None, type=int,
                        help='When to start the exponential moving average of G (default: halflife_k)')
    parser.add_argument('--halflife_lr', default=0, type=int, help='Apply LR decay when > 0')

    parser.add_argument('--use_nerf_proj', action='store_true')

    # Options for logging specification
    parser.add_argument('--no_fid', action='store_true',
                        help='Do not track FIDs during training')
    parser.add_argument('--no_gif', action='store_true',
                        help='Do not save GIF of sample generations from a fixed latent periodically during training')
    parser.add_argument('--n_eval_avg', default=3, type=int,
                        help='How many times to average FID and IS')
    parser.add_argument('--print_every', help='', default=1000, type=int)
    parser.add_argument('--evaluate_every', help='', default=2000, type=int)
    parser.add_argument('--save_every', help='', default=10000, type=int)
    parser.add_argument('--comment', help='Comment', default='', type=str)

    # Options for resuming / fine-tuning
    parser.add_argument('--resume', default=None, type=str,
                        help='Path to logdir to resume the training')
    parser.add_argument('--finetune', default=None, type=str,
                        help='Path to logdir that contains a pre-trained checkpoint of D')

    return parser.parse_args()

def _update_warmup(optimizer, cur_step, warmup, lr):
    if warmup > 0:
        ratio = min(1., (cur_step + 1) / (warmup + 1e-8))
        lr_w = ratio * lr
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr_w


def _update_lr(optimizer, cur_step, batch_size, halflife_lr, lr, mult=1.0):
    if halflife_lr > 0 and (cur_step > 0) and (cur_step % 10000 == 0):
        #ratio = (cur_step * batch_size) / halflife_lr
        ratio=cur_step/10000-11
        lr_mul = 0.5 ** ratio
        lr_w = lr_mul * lr * mult
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr_w
        return lr_w
    return None


def r1_loss(D, images, augment_fn):
    images_aug = augment_fn(images).detach()
    images_aug.requires_grad = True
    d_real = D(images_aug)
    grad_real, = autograd.grad(outputs=d_real.sum(), inputs=images_aug,
                               create_graph=True, retain_graph=True)
    grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()
    return grad_penalty


def _sample_generator(G, num_samples, enable_grad=True,imgs=None,illus=None):
    latent_samples = G.sample_latent(num_samples)
    if enable_grad:
        generated_data = G(x=imgs,input=latent_samples,illu=illus)
        #print(summary(G, imgs.shape,latent_samples.shape))
    else:
        with torch.no_grad():
            generated_data = G(x=imgs,input=latent_samples,illu=illus)
            #print(summary(G, imgs.shape,latent_samples.shape))
    return generated_data

@gin.configurable("options")
def get_options_dict(dataset=gin.REQUIRED,
                     loss=gin.REQUIRED,
                     batch_size=32, fid_size=10000,
                     max_steps=200, warmup=0, n_critic=1,
                     lr=0.002, lr_d=None, beta=(.0, .99),
                     lbd=1., lbd2=1.):
    if lr_d is None:
        lr_d = lr
    return {
        "dataset": dataset,
        "batch_size": batch_size,
        "fid_size": fid_size,
        "loss": loss,
        "max_steps": max_steps, "warmup": warmup,
        "n_critic": n_critic,
        "lr": lr, "lr_d": lr_d, "beta": beta,
        "lbd": lbd, "lbd2": lbd2
    }


def train(P, opt, train_fn, models, optimizers, 
            ltrain_ltarget_pair_loader,
            presudo_pair_loader,
            val_pair_loader,logger):
    generator, discriminator_single, g_ema = models
    opt_G,opt_DS = optimizers
    for param_group in opt_G.param_groups:
        param_group['lr'] = opt["lr"]
        print(f"G_lr为{param_group['lr']}")
    for param_group in opt_DS.param_groups:
        param_group['lr'] = opt["lr"]
        print(f"DS_lr为{param_group['lr']}")
    
    losses = {'G_loss': [],'G_critic_loss': [], 'G_l_mse_loss': [], 'G_ch_loss': [],
                'DS_loss': [], 'DS_penalty': [],'DS_real': [], 'DS_gen': [], 'DS_r1': []
              }
    metrics={'fid_score':[]}
    # metrics={}
    # metrics['fid_score'] = FIDScore(opt['dataset'], opt['fid_size'], P.n_eval_avg)
    metric_SSIM = SSIM(data_range=1.0)
    metric_PSNR = PSNR(data_range=1.0)
    #metric_FID = FID(num_features=192, feature_extractor=default_model)
    metric_SSIM.attach(default_evaluator, 'ssim')
    metric_PSNR.attach(default_evaluator, 'psnr')
    #metric_FID.attach(default_evaluator, 'fid')
    logger.log_dirname("Steps {}".format(P.starting_step))

    for step in range(P.starting_step, opt['max_steps'] + 1):
        if step % P.evaluate_every == 0:
            val_images,val_images128,val_target_images,val_illus=next(val_pair_loader)
            val_images = val_images.cuda()
            val_target_images=val_target_images.cuda()
            val_illus=val_illus.cuda()
            val_images128=val_images128.cuda()
            with torch.no_grad():
                val_gen_images = _sample_generator(generator, val_images.size(0),enable_grad=True,imgs=val_images,illus=val_illus)
                val_ch_loss=ch_loss(val_gen_images,img_size=128)
                val_target_loss=F.smooth_l1_loss(val_gen_images,val_target_images)
            #writer.add_scalar('stage2_val_fid',val_fid_value,step)
            #print(f"图片类型为{val_target_images.dtype}和{val_gen_images.dtype}")torch.float32
            #print(f"{torch.max(val_target_images)-torch.min(val_target_images)}")1.0
            #print(f"{torch.max(val_gen_images)-torch.min(val_gen_images)}")1.0
            state = default_evaluator.run([[val_gen_images,val_target_images]])
            fid_value=my_fid_score(path_base='base_stats.npz', G=generator, size=val_images.size(0), batch_size=val_images.size(0), model=None, dims=192)
            metrics['fid_score'].append(fid_value)
            writer.add_scalar('stage2_val_SSIM',state.metrics['ssim'],step)
            writer.add_scalar('stage2_val_PSNR',state.metrics['psnr'],step)
            writer.add_scalar('stage2_val_FID',fid_value,step)
            #writer.add_scalar('stage2_val_FID',state.metrics['fid'],step)
            writer.add_scalar('stage2_val_ch_loss',val_ch_loss,step)
            writer.add_scalar('stage2_val_target_loss',val_target_loss,step)
            # logger.log('[Steps %7d][stage2_val_SSIM %.7f][stage2_val_PSNR %.7f][stage2_val_FID %.7f][stage2_val_ch_loss %.7f] [stage2_val_target_loss %.14f]' %
            #     (step, state.metrics['ssim'], state.metrics['psnr'],state.metrics['fid'], val_ch_loss, val_target_loss))
            logger.log('[Steps %7d][stage2_val_SSIM %.7f][stage2_val_PSNR %.7f][stage2_val_ch_loss %.7f] [stage2_val_target_loss %.14f]' %
                (step, state.metrics['ssim'], state.metrics['psnr'],val_ch_loss, val_target_loss))
        d_regularize = (step % P.d_reg_every == 0) and (P.lbd_r1 > 0)

        if P.use_warmup:
            _update_warmup(opt_G, step, opt["warmup"], opt["lr"])
            _update_warmup(opt_DS, step, opt["warmup"], opt["lr_d"])
        if (not P.use_warmup) or step > opt["warmup"]:
            cur_lr_g = _update_lr(opt_G, step, opt["batch_size"], P.halflife_lr, opt["lr"])
            cur_lr_ds = _update_lr(opt_DS, step, opt["batch_size"], P.halflife_lr, opt["lr_d"])
            if cur_lr_ds and cur_lr_g:
                logger.log('LR Updated: [G %.10f][DS %.10f]' % (cur_lr_g,cur_lr_ds))
        do_ema = (step * opt['batch_size']) > (P.ema_start_k * 1000)
        accum = P.accum if do_ema else 0
        accumulate(g_ema, generator, accum)

        # Start discriminator training
        generator.train()
        discriminator_single.train()
        images,target_images,illus,real_images=next(presudo_pair_loader)
        images=images.cuda()
        target_images=target_images.cuda()
        illus=illus.cuda()
        real_images=real_images.cuda()
        ltrain_images, ltarget_images,lillus,lgan_images = next(ltrain_ltarget_pair_loader)
        ltrain_images=ltrain_images.cuda()
        ltarget_images=ltarget_images.cuda()
        lillus=lillus.cuda()
        lgan_images=lgan_images.cuda()
        set_grad(generator, False)
        set_grad(discriminator_single, True)
        ugen_images = _sample_generator(generator, images.size(0),enable_grad=True,imgs=images,illus=illus)

        #ds
        ds_loss, ds_aux = train_fn["train3_D_match"](P, discriminator_single, opt,real_images,ugen_images)
        loss = ds_loss+ ds_aux['penalty']

        opt_DS.zero_grad()
        loss.backward()

        opt_DS.step()

        losses['DS_loss'].append(ds_loss.item())
        losses['DS_real'].append(ds_aux['d_real'].item())
        losses['DS_gen'].append(ds_aux['d_gen'].item())
        losses['DS_penalty'].append(ds_aux['penalty'].item())
        writer.add_scalars('stage2_D',{'ds_loss': losses['DS_loss'][-1], 'ds_penalty': losses['DS_penalty'][-1]}, step)




        # Start generator training
        set_grad(generator, True)
        set_grad(discriminator_single, False)
        ugen_images = _sample_generator(generator, images.size(0),enable_grad=True,imgs=images,illus=illus)
        gs_loss ,gs_aux= train_fn["train3_G_match"](P, discriminator_single, opt,ugen_images,target_images,ugen_images)

        g_loss=gs_loss
        opt_G.zero_grad()
        g_loss.backward()
        opt_G.step()


    
        losses['G_loss'].append(g_loss.item())
        losses['G_critic_loss'].append(gs_aux['critic_loss'].item())
        losses['G_l_mse_loss'].append(gs_aux['l_mse_loss'].item())
        losses['G_ch_loss'].append(gs_aux['g_ch_loss'].item())
      

        writer.add_scalars('stage2_G_D',{'g_loss': g_loss.item(), 'd_loss': loss.item()}, step)
        writer.add_scalars('stage2_G',{'g_critic_loss': losses['G_critic_loss'][-1], 'g_l_mse_loss': losses['G_l_mse_loss'][-1],'g_ch_loss':losses['G_ch_loss'][-1]}, step)

        generator.eval()
        discriminator_single.eval()

        if step % P.print_every == 0:
            logger.log('[Steps %7d][G %.7f][G_critic %.7f] [G_l_mse %.7f] [G_ch_loss %.14f][DS %.7f][DS_real %.7f][DS_gen %.7f][DS_penalty %.7f]' %
                       (step, losses['G_loss'][-1], losses['G_critic_loss'][-1], losses['G_l_mse_loss'][-1], losses['G_ch_loss'][-1],
                        losses['DS_loss'][-1],losses['DS_real'][-1], losses['DS_gen'][-1],losses['DS_penalty'][-1]))

            for name in losses:
                values = losses[name]
                if len(values) > 0:
                    logger.scalar_summary('gan/train/' + name, values[-1], step)
    

        if step % P.evaluate_every == 0:
            logger.log_dirname("Steps {}".format(step + 1))
            #wandb.log({"augmented_real_images": wandb.Image(aug_grid), "generated_images": wandb.Image(fixed_gen.summary()[-1])}, step=step)
            #fid_score = metrics.get('fid_score')
            G_state_dict = generator.module.state_dict()
            #DP_state_dict = discriminator_pair.module.state_dict()
            DS_state_dict = discriminator_single.module.state_dict()
            Ge_state_dict = g_ema.state_dict()
            #fid_value=my_fid_score(path_base='base_stats.npz', G=generator, size=images.size(0), batch_size=images.size(0), model=None, dims=192)
            #metrics['fid_score'].append(fid_value)
            logger.log('[Steps %7d][fid_score %.7f]' %(step, metrics['fid_score'][-1]))



            torch.save(G_state_dict, logger.logdir + '/gen_stage3_Ds.pt')
            torch.save(DS_state_dict, logger.logdir + '/disS_stage3_Ds.pt')
            torch.save(Ge_state_dict, logger.logdir + '/gen_ema_stage3_Ds.pt')
            # if fid_score and fid_score.is_best:
            #     torch.save(G_state_dict, logger.logdir + '/gen_best.pt')
            #     torch.save(DS_state_dict, logger.logdir + '/dis_best.pt')
            #     torch.save(Ge_state_dict, logger.logdir + '/gen_ema_best.pt')
            if step % P.save_every == 0:
                torch.save(G_state_dict, logger.logdir + f'/gen_stage3_Ds_{step}.pt')
                torch.save(DS_state_dict, logger.logdir + f'/disS_stage3_Ds_{step}.pt')
                torch.save(Ge_state_dict, logger.logdir + f'/gen_ema_stage3_Ds_{step}.pt')
                torch.save({'epoch': step,'optim_G': opt_G.state_dict(),'optim_DS': opt_DS.state_dict(),
            }, logger.logdir + f'/optim_stage3_Ds_{step}.pt')
            
            torch.save({
                'epoch': step,
                'optim_G': opt_G.state_dict(),
                'optim_DS': opt_DS.state_dict(),
            }, logger.logdir + '/optim_stage3_Ds.pt')


def worker(P):
    gin.parse_config_files_and_bindings(['configs/defaults/gan.gin',
                                         'configs/defaults/augment.gin',
                                         P.gin_config], [])

    options = get_options_dict()
    seed=10
    torch.manual_seed(seed)
    ltrain_lgan_ltarget_pair_set,image_size=get_dataset(dataset='labeled_data_stage3')
    presudo_pair_set,resolution= get_dataset(dataset='unlabeled_data1_LAB_presudo_stage3')
    val_pair_set,val_resolution=get_dataset(dataset='val_data_stage3')

    seed=10
    torch.manual_seed(seed)
    ltrain_ltarget_pair_loader=DataLoader(ltrain_lgan_ltarget_pair_set, shuffle=True, pin_memory=True, num_workers=P.workers,
                              batch_size=options['batch_size'], drop_last=True)
    presudo_pair_loader=DataLoader(presudo_pair_set, shuffle=True, pin_memory=True, num_workers=P.workers,
                              batch_size=options['batch_size'], drop_last=True)
    val_pair_loader = DataLoader(val_pair_set, shuffle=False, pin_memory=False, num_workers=P.workers,
                              batch_size=50, drop_last=False)    
    ltrain_ltarget_pair_loader = cycle4(ltrain_ltarget_pair_loader)
    presudo_pair_loader=cycle4(presudo_pair_loader)
    val_pair_loader=cycle4(val_pair_loader)


    if P.no_lazy:
        P.d_reg_every = 1
    if P.ema_start_k is None:
        P.ema_start_k = P.halflife_k
        

    P.accum = 0.5 ** (options['batch_size'] / (P.halflife_k * 1000))
    from vit_generator_skip import vit_my_8
    resolution = image_size[0]
    generator = vit_my_8(patch_size=16)
    g_ema = vit_my_8(patch_size=16)
    discriminator_pair = ResidualDiscriminatorP(size=resolution, small32=False,mlp_linear=True, d_hidden=512)
    discriminator_single = ResidualDiscriminatorP(size=resolution, small32=False,mlp_linear=True, d_hidden=512,input_channel=3)
    if P.resume:
        print(f"=> Loading checkpoint from '{P.resume}'")
        state_G = torch.load(f"{P.resume}/gen_stage2.pt")
        #state_DP = torch.load(f"{P.resume}/disP_stage3.pt")
        #state_DS = torch.load(f"{P.resume}/disS_stage3.pt")
        state_Ge = torch.load(f"{P.resume}/gen_ema_stage2.pt")
        # state_G = torch.load(f"{P.resume}/gen_50000_stage2.pt")
        # state_Ge = torch.load(f"{P.resume}/gen_ema_50000_stage2.pt")
        #state_G = torch.load(f"{P.resume}/gen_best.pt")
        #print(f"state_G为{state_G.items()}")
        generator.load_state_dict(state_G,strict=False)
        #discriminator_pair.load_state_dict(state_DP,strict=True)
        #discriminator_single.load_state_dict(state_DS,strict=True)
        g_ema.load_state_dict(state_Ge,strict=False)
        
    if P.finetune:
        print(f"=> Loading checkpoint for fine-tuning: '{P.finetune}'")
        #state_DP = torch.load(f"{P.finetune}/dis_pair.pt")
        #discriminator_pair.load_state_dict(state_DP, strict=False)
        #discriminator_pair.reset_parameters(discriminator_pair.linear)
        state_DS = torch.load(f"{P.finetune}/dis_single.pt")
        discriminator_single.load_state_dict(state_DS, strict=False)
        discriminator_single.reset_parameters(discriminator_single.linear)
        P.comment += 'ft'

    generator = generator.cuda()
    #discriminator_pair = discriminator_pair.cuda()
    discriminator_single = discriminator_single.cuda()
    g_ema = g_ema.cuda()
    g_ema.eval()

    for name, param in generator.named_parameters():
        if "cls_token" in name or "pos_embed" in name or "style." in name or "blocks." in name or "norm." in name:
            param.requires_grad=False 
    G_optimizer = optim.Adam(filter(lambda p: p.requires_grad, generator.parameters()),lr=options["lr"], betas=options["beta"]) 
    D_optimizer_single = optim.Adam(discriminator_single.parameters(),
                             lr=options["lr_d"], betas=options["beta"])

    if P.resume:
        logger = Logger(None, resume=P.resume)
        #wandb.init(project='vitgan', name=f'{P.gin_stem}_{P.architecture}_' + f'{P.filename}_{_desc}{P.comment}', resume=True)

        #wandb.config.update(P)
        #wandb.config.update(options)

        #wandb.watch(generator)
        #wandb.watch(discriminator)
    else:
        _desc = f"R{P.lbd_r1}_H{P.halflife_k}"
        if P.halflife_lr > 0:
            _desc += f"_lr{P.halflife_lr / 1000000:.1f}M"
        _desc += f"_NoLazy" if P.no_lazy else "_Lazy"
        _desc += f"_Warmup" if P.use_warmup else "_NoWarmup"

        logger = Logger(f'{P.filename}_{_desc}{P.comment}', subdir=f'gan_dp/{P.gin_stem}/{P.architecture}')
        #wandb.init(project='vitgan', name=f'{P.gin_stem}_{P.architecture}_' + f'{P.filename}_{_desc}{P.comment}')
        
        #wandb.config.update(P)
        #wandb.config.update(options)

        #wandb.watch(generator)
        #wandb.watch(discriminator)

        shutil.copy2(P.gin_config, f"{logger.logdir}/config.gin")
    P.logdir = logger.logdir
    P.eval_seed = np.random.randint(10000)

    if P.resume:
        opt = torch.load(f"{P.resume}/optim_stage2.pt")
        # G_optimizer.load_state_dict(opt['optim_G'])
        # D_optimizer_single.load_state_dict(opt['optim_DS'])
        logger.log(f"Checkpoint loaded from '{P.resume}'")
        P.starting_step = opt['epoch'] + 1
    else:
        logger.log(generator)
        logger.log(discriminator_pair)
        logger.log(discriminator_single)
        logger.log(f"# Params - G: {count_parameters(generator)}, D_single: {count_parameters(discriminator_single)}")
        logger.log(options)
        P.starting_step = 1
    logger.log(f"Use G moving average: {P.accum}")

    if P.finetune:
        logger.log(f"Checkpoint loaded from '{P.finetune}'")



    P.augment_fn = get_augment(mode=P.aug).cuda()
    generator = nn.DataParallel(generator)
    g_ema = nn.DataParallel(g_ema)
    generator.sample_latent = generator.module.sample_latent
    discriminator_single = nn.DataParallel(discriminator_single)


    train(P, options, P.train_fn,
          models=(generator, discriminator_single, g_ema),
          optimizers=(G_optimizer,D_optimizer_single),
            ltrain_ltarget_pair_loader=ltrain_ltarget_pair_loader,
            presudo_pair_loader=presudo_pair_loader,
            val_pair_loader=val_pair_loader,logger=logger)


if __name__ == '__main__':
    P = parse_args()
    if P.comment:
        P.comment = '_' + P.comment
    P.gin_stem = Path(P.gin_config).stem
    P = setup(P)
    P.distributed = False
    worker(P)