text: import argparse
  import os
  import random
  import yaml
  import time
  import logging
  import pprint
  
  import scipy.stats as stats
  import torch.backends.cudnn as cudnn
  import torch.optim as optim
  import torchvision.utils as vutils
  import numpy as np
  from tensorboardX import SummaryWriter
  from torch.autograd import grad
  from easydict import EasyDict
  
  from data.train import CreateDataLoader as train_loader
  from data.eval import CreateDataLoader as val_loader
  from data.dataloader import get_dataloader
  from utils import create_logger, save_checkpoint, load_state, get_scheduler, AverageMeter, calculate_fid
  from models.standard import *
  
  parser = argparse.ArgumentParser(description='PyTorch Colorization Training')
  
  parser.add_argument('--config', default='experiments/origin/config.yaml')
  parser.add_argument('--resume', default='', type=str, help='path to checkpoint')
  
  batch_size = 16
  def calc_gradient_penalty(netD, real_data, fake_data, sketch_feat):
      alpha = torch.rand(batch_size, 1, 1, 1, device=config.device)
  
      interpolates = alpha * real_data + ((1 - alpha) * fake_data)
  
      interpolates.requires_grad = True
  
      disc_interpolates = netD(interpolates, sketch_feat)
  
      gradients = grad(outputs=disc_interpolates, inputs=interpolates,
                       grad_outputs=torch.ones(disc_interpolates.size(), device=config.device), create_graph=True,
                       retain_graph=True, only_inputs=True)[0]
  
      gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * config.gpW
      return gradient_penalty
  
  
  def mask_gen():
      maskS =256 // 4
  
      mask1 = torch.cat(
          [torch.rand(1, 1, maskS, maskS).ge(X.rvs(1)[0]).float() for _ in range(batch_size // 2)], 0)
      mask2 = torch.cat([torch.zeros(1, 1, maskS, maskS).float() for _ in range(batch_size // 2)], 0)
      mask = torch.cat([mask1, mask2], 0)
  
      return mask.to(config.device)
  
  
  def main():
      global args, config, X
  
      args = parser.parse_args()
      print(args)
  
      with open(args.config) as f:
          config = EasyDict(yaml.load(f))
  
      config.save_path = os.path.dirname(args.config)
  
      ####### regular set up
      assert torch.cuda.is_available()
      device = torch.device("cuda")
      config.device = device
  
      # random seed setup
      print("Random Seed: ", config.seed)
      random.seed(config.seed)
      torch.manual_seed(config.seed)
      torch.cuda.manual_seed(config.seed)
      cudnn.benchmark = True
  
      ####### regular set up end
  
      netG = NetG(ngf=config.ngf)
      netD = NetD(ndf=config.ndf)
  
      netF = NetF()
      netI = NetI().eval()
  
      # netG = torch.nn.DataParallel(NetG(ngf=config.ngf))
      # netD = torch.nn.DataParallel(NetD(ndf=config.ndf))
  
      # netF = torch.nn.DataParallel(NetF())
      # netI = torch.nn.DataParallel(NetI()).eval()
      for param in netF.parameters():
          param.requires_grad = False
  
      criterion_MSE = nn.MSELoss()
  
      fixed_sketch = torch.tensor(0, device=device).float()
      fixed_hint = torch.tensor(0, device=device).float()
      fixed_sketch_feat = torch.tensor(0, device=device).float()
  
      ####################
      netD = netD.to(device)
      netG = netG.to(device)
      netF = netF.to(device)
      netI = netI.to(device)
      criterion_MSE = criterion_MSE.to(device)
  
      # setup optimizer
  
      optimizerG = optim.Adam(netG.parameters(), lr=config.lr_scheduler.base_lr, betas=(0.5, 0.9))
      optimizerD = optim.Adam(netD.parameters(), lr=config.lr_scheduler.base_lr, betas=(0.5, 0.9))
  
      last_iter = -1
      best_fid = 1e6
  
      if args.resume:
          best_fid, last_iter = load_state(args.resume, netG, netD, optimizerG, optimizerD)
  
      config.lr_scheduler['last_iter'] = last_iter
  
      config.lr_scheduler['optimizer'] = optimizerG
      lr_schedulerG = get_scheduler(config.lr_scheduler)
      config.lr_scheduler['optimizer'] = optimizerD
      lr_schedulerD = get_scheduler(config.lr_scheduler)
  
      tb_logger = SummaryWriter(config.save_path + '/events')
      logger = create_logger('global_logger', config.save_path + '/log.txt')
      logger.info(f'args: {pprint.pformat(args)}')
      logger.info(f'config: {pprint.pformat(config)}')
  
      batch_time = AverageMeter(config.print_freq)
      data_time = AverageMeter(config.print_freq)
      flag = 1
      mu, sigma = 1, 0.005
      X = stats.truncnorm((0 - mu) / sigma, (1 - mu) / sigma, loc=mu, scale=sigma)
      i = 0
      curr_iter = last_iter + 1
  
      #dataloader = train_loader(config)
      dataloader = get_dataloader('yumi', 'train', batch_size=batch_size)
      data_iter = iter(dataloader)
  
      end = time.time()
      while i < len(dataloader):
          lr_schedulerG.step(curr_iter)
          lr_schedulerD.step(curr_iter)
          current_lr = lr_schedulerG.get_lr()[0]
          ############################
          # (1) Update D network
          ###########################
          for p in netD.parameters():  # reset requires_grad
              p.requires_grad = True  # they are set to False below in netG update
          for p in netG.parameters():
              p.requires_grad = False  # to avoid computation ft_params
  
          # train the discriminator Diters times
          j = 0
          while j < config.diters:
              netD.zero_grad()
  
              i += 1
              j += 1
  
              data_end = time.time()
              real_cim, real_vim, real_sim = data_iter.next()
              data_time.update(time.time() - data_end)
  
              real_cim, real_vim, real_sim = real_cim.to(device), real_vim.to(device), real_sim.to(device)
              mask = mask_gen()
             
              hint = torch.cat((real_vim * mask, mask), 1)
  
              # train with fake
              with torch.no_grad():
                  feat_sim = netI(real_sim).detach()
                  fake_cim = netG(real_sim, hint, feat_sim).detach()
  
              errD_fake = netD(fake_cim, feat_sim)
              errD_fake = errD_fake.mean(0).view(1)
  
              errD_fake.backward(retain_graph=True)  # backward on score on real
  
              errD_real = netD(real_cim, feat_sim)
              errD_real = errD_real.mean(0).view(1)
              errD = errD_real - errD_fake
  
              errD_realer = -1 * errD_real + errD_real.pow(2) * config.drift
  
              errD_realer.backward(retain_graph=True)  # backward on score on real
  
              gradient_penalty = calc_gradient_penalty(netD, real_cim, fake_cim, feat_sim)
              gradient_penalty.backward()
  
              optimizerD.step()
  
          ############################
          # (2) Update G network
          ############################
  
          for p in netD.parameters():
              p.requires_grad = False  # to avoid computation
          for p in netG.parameters():
              p.requires_grad = True
          netG.zero_grad()
  
          data = data_iter.next()
          real_cim, real_vim, real_sim = data
          i += 1
  
          real_cim, real_vim, real_sim = real_cim.to(device), real_vim.to(device), real_sim.to(device)
  
          if flag:  # fix samples
              mask = mask_gen()
              hint = torch.cat((real_vim * mask, mask), 1)
              with torch.no_grad():
                  feat_sim = netI(real_sim).detach()
  
              tb_logger.add_image('target imgs', vutils.make_grid(real_cim.mul(0.5).add(0.5), nrow=4))
              tb_logger.add_image('sketch imgs', vutils.make_grid(real_sim.mul(0.5).add(0.5), nrow=4))
              tb_logger.add_image('hint', vutils.make_grid((real_vim * mask).mul(0.5).add(0.5), nrow=4))
  
              fixed_sketch.resize_as_(real_sim).copy_(real_sim)
              fixed_hint.resize_as_(hint).copy_(hint)
              fixed_sketch_feat.resize_as_(feat_sim).copy_(feat_sim)
  
              flag -= 1
  
          mask = mask_gen()
          hint = torch.cat((real_vim * mask, mask), 1)
  
          with torch.no_grad():
              feat_sim = netI(real_sim).detach()
  
          fake = netG(real_sim, hint, feat_sim)
  
          errd = netD(fake, feat_sim)
          errG = errd.mean() * config.advW * -1
          errG.backward(retain_graph=True)
          feat1 = netF(fake)
          with torch.no_grad():
              feat2 = netF(real_cim)
  
          contentLoss = criterion_MSE(feat1, feat2)
          contentLoss.backward()
  
          optimizerG.step()
          batch_time.update(time.time() - end)
  
          ############################
          # (3) Report & 100 Batch checkpoint
          ############################
          curr_iter += 1
  
          if curr_iter % config.print_freq == 0:
              tb_logger.add_scalar('VGG MSE Loss', contentLoss.item(), curr_iter)
              tb_logger.add_scalar('wasserstein distance', errD.item(), curr_iter)
              tb_logger.add_scalar('errD_real', errD_real.item(), curr_iter)
              tb_logger.add_scalar('errD_fake', errD_fake.item(), curr_iter)
              tb_logger.add_scalar('Gnet loss toward real', errG.item(), curr_iter)
              tb_logger.add_scalar('gradient_penalty', gradient_penalty.item(), curr_iter)
              tb_logger.add_scalar('lr', current_lr, curr_iter)
              logger.info(f'Iter: [{curr_iter}/{len(dataloader)//(config.diters+1)}]\t'
                          f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                          f'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                          f'errG {errG.item():.4f}\t'
                          f'errD {errD.item():.4f}\t'
                          f'err_D_real {errD_real.item():.4f}\t'
                          f'err_D_fake {errD_fake.item():.4f}\t'
                          f'content loss {contentLoss.item():.4f}\t'
                          f'LR {current_lr:.4f}'
                          )
  
          if curr_iter % config.print_img_freq == 0:
              with torch.no_grad():
                  fake = netG(fixed_sketch, fixed_hint, fixed_sketch_feat)
                  tb_logger.add_image('colored imgs',
                                      vutils.make_grid(fake.detach().mul(0.5).add(0.5), nrow=4),
                                      curr_iter)
  
          if curr_iter % config.val_freq == 0:
              fid, var = validate(netG, netI)
              tb_logger.add_scalar('fid_val', fid, curr_iter)
              tb_logger.add_scalar('fid_variance', var, curr_iter)
              logger.info(f'fid: {fid:.3f} ({var})\t')
  
              # remember best fid and save checkpoint
              is_best = fid < best_fid
              best_fid = min(fid, best_fid)
              save_checkpoint({
                  'step': curr_iter - 1,
                  'state_dictG': netG.state_dict(),
                  'state_dictD': netD.state_dict(),
                  'state_dictI': netI.state_dict(),
                  'best_fid': best_fid,
                  'optimizerG': optimizerG.state_dict(),
                  'optimizerD': optimizerD.state_dict(),
              }, is_best, config.save_path + '/ckpt')
  
          end = time.time()
  
          # if curr_iter == 200:
          #     print("Epoch 200 FInish!!!!!!!!!")
          #     break
  
  
  def validate(netG, netI):
      fids = []
      fid_value = 0
      for _ in range(3):
          fid = calculate_fid(netG, netI, get_dataloader('yumi', 'val', batch_size=batch_size), config, 2048)
          print('FID: ', fid)
          fid_value += fid
          fids.append(fid)
      fid_value /= 3
      return fid_value, np.var(fids)
  
  if __name__ == '__main__':
      main()
  