code/studiogan/src/worker.py

# PyTorch StudioGAN: https://github.com/POSTECH-CVLab/PyTorch-StudioGAN
# The MIT License (MIT)
# See license file or visit https://github.com/POSTECH-CVLab/PyTorch-StudioGAN for details

# src/worker.py

from os.path import join
import os
import sys
import glob
import random
import string
import pickle
import copy

from torch.nn import DataParallel
from torch.nn.parallel import DistributedDataParallel
from torchvision import transforms
from PIL import Image
from tqdm import tqdm
from scipy import ndimage
from utils.style_ops import conv2d_gradfix
from utils.style_ops import upfirdn2d
from sklearn.manifold import TSNE
from datetime import datetime
import torch
import torchvision
import torch.nn as nn
import torch.distributed as dist
import torch.nn.functional as F
import numpy as np

import metrics.features as features
import metrics.ins as ins
import metrics.fid as fid
import metrics.prdc as prdc
import metrics.resnet as resnet
import utils.ckpt as ckpt
import utils.sample as sample
import utils.misc as misc
import utils.losses as losses
import utils.sefa as sefa
import utils.ops as ops
import utils.resize as resize
import utils.apa_aug as apa_aug
import wandb

SAVE_FORMAT = "step={step:0>3}-Inception_mean={Inception_mean:<.4}-Inception_std={Inception_std:<.4}-FID={FID:<.5}.pth"

LOG_FORMAT = ("Step: {step:>6} "
              "Progress: {progress:<.1%} "
              "Elapsed: {elapsed} "
              "Gen_loss: {gen_loss:<.4} "
              "Dis_loss: {dis_loss:<.4} "
              "Cls_loss: {cls_loss:<.4} "
              "Topk: {topk:>4} "
              "aa_p: {aa_p:<.4} ")


class WORKER(object):
    def __init__(self, cfgs, run_name, Gen, Gen_mapping, Gen_synthesis, Dis, Gen_ema, Gen_ema_mapping, Gen_ema_synthesis,
                 ema, eval_model, train_dataloader, eval_dataloader, global_rank, local_rank, mu, sigma, real_feats, logger,
                 aa_p, best_step, best_fid, best_ckpt_path, lecam_emas, num_eval, loss_list_dict, metric_dict_during_train):
        self.cfgs = cfgs
        self.run_name = run_name
        self.Gen = Gen
        self.Gen_mapping = Gen_mapping
        self.Gen_synthesis = Gen_synthesis
        self.Dis = Dis
        self.Gen_ema = Gen_ema
        self.Gen_ema_mapping = Gen_ema_mapping
        self.Gen_ema_synthesis = Gen_ema_synthesis
        self.ema = ema
        self.eval_model = eval_model
        self.train_dataloader = train_dataloader
        self.eval_dataloader = eval_dataloader
        self.global_rank = global_rank
        self.local_rank = local_rank
        self.mu = mu
        self.sigma = sigma
        self.real_feats = real_feats
        self.logger = logger
        self.aa_p = aa_p
        self.best_step = best_step
        self.best_fid = best_fid
        self.best_ckpt_path = best_ckpt_path
        self.lecam_emas = lecam_emas
        self.num_eval = num_eval
        self.loss_list_dict = loss_list_dict
        self.metric_dict_during_train = metric_dict_during_train
        self.metric_dict_during_final_eval = {}

        self.cfgs.define_augments(local_rank)
        self.cfgs.define_losses()
        self.DATA = cfgs.DATA
        self.MODEL = cfgs.MODEL
        self.LOSS = cfgs.LOSS
        self.STYLEGAN = cfgs.STYLEGAN
        self.OPTIMIZATION = cfgs.OPTIMIZATION
        self.PRE = cfgs.PRE
        self.AUG = cfgs.AUG
        self.RUN = cfgs.RUN
        self.MISC = cfgs.MISC
        self.is_stylegan = cfgs.MODEL.backbone in ["stylegan2", "stylegan3"]
        self.effective_batch_size = self.OPTIMIZATION.batch_size * self.OPTIMIZATION.acml_steps
        self.blur_init_sigma = self.STYLEGAN.blur_init_sigma
        self.blur_fade_kimg = self.effective_batch_size * 200/32
        self.DDP = self.RUN.distributed_data_parallel
        self.adc_fake = False

        num_classes = self.DATA.num_classes

        self.sampler = misc.define_sampler(self.DATA.name, self.MODEL.d_cond_mtd,
                                            self.OPTIMIZATION.batch_size, self.DATA.num_classes)

        self.pl_reg = losses.PathLengthRegularizer(device=local_rank, pl_weight=cfgs.STYLEGAN.pl_weight, pl_no_weight_grad=(cfgs.MODEL.backbone == "stylegan2"))
        self.l2_loss = torch.nn.MSELoss()
        self.ce_loss = torch.nn.CrossEntropyLoss()
        self.fm_loss = losses.feature_matching_loss
        self.lecam_ema = ops.LeCamEMA()
        if self.lecam_emas is not None:
            self.lecam_ema.__dict__ = self.lecam_emas
        self.lecam_ema.decay, self.lecam_ema.start_itr = self.LOSS.lecam_ema_decay, self.LOSS.lecam_ema_start_iter
        if self.LOSS.adv_loss == "MH":
            self.lossy = torch.LongTensor(self.OPTIMIZATION.batch_size).to(self.local_rank)
            self.lossy.data.fill_(self.DATA.num_classes)

        if self.AUG.apply_ada + self.AUG.apply_apa:
            if self.AUG.apply_ada: self.AUG.series_augment.p.copy_(torch.as_tensor(self.aa_p))
            self.aa_interval = self.AUG.ada_interval if self.AUG.ada_interval != "N/A" else self.AUG.apa_interval
            self.aa_target = self.AUG.ada_target if self.AUG.ada_target != "N/A" else self.AUG.apa_target
            self.aa_kimg = self.AUG.ada_kimg if self.AUG.ada_kimg != "N/A" else self.AUG.apa_kimg
            self.dis_sign_real, self.dis_sign_fake = torch.zeros(2, device=self.local_rank), torch.zeros(2, device=self.local_rank)
            self.dis_logit_real, self.dis_logit_fake = torch.zeros(2, device=self.local_rank), torch.zeros(2, device=self.local_rank)
            self.dis_sign_real_log, self.dis_sign_fake_log = torch.zeros(2, device=self.local_rank), torch.zeros(2, device=self.local_rank)
            self.dis_logit_real_log, self.dis_logit_fake_log = torch.zeros(2, device=self.local_rank), torch.zeros(2, device=self.local_rank)

        if self.MODEL.aux_cls_type == "ADC":
            num_classes = num_classes*2
            self.adc_fake = True

        if self.MODEL.d_cond_mtd == "AC":
            self.cond_loss = losses.CrossEntropyLoss()
        elif self.MODEL.d_cond_mtd == "2C":
            self.cond_loss = losses.ConditionalContrastiveLoss(num_classes=num_classes,
                                                               temperature=self.LOSS.temperature,
                                                               master_rank="cuda",
                                                               DDP=self.DDP)
        elif self.MODEL.d_cond_mtd == "D2DCE":
            self.cond_loss = losses.Data2DataCrossEntropyLoss(num_classes=num_classes,
                                                              temperature=self.LOSS.temperature,
                                                              m_p=self.LOSS.m_p,
                                                              master_rank="cuda",
                                                              DDP=self.DDP)
        else: pass

        if self.MODEL.aux_cls_type == "TAC":
            self.cond_loss_mi = copy.deepcopy(self.cond_loss)

        self.gen_ctlr = misc.GeneratorController(generator=self.Gen_ema if self.MODEL.apply_g_ema else self.Gen,
                                                 generator_mapping=self.Gen_ema_mapping,
                                                 generator_synthesis=self.Gen_ema_synthesis,
                                                 batch_statistics=self.RUN.batch_statistics,
                                                 standing_statistics=False,
                                                 standing_max_batch="N/A",
                                                 standing_step="N/A",
                                                 cfgs=self.cfgs,
                                                 device=self.local_rank,
                                                 global_rank=self.global_rank,
                                                 logger=self.logger,
                                                 std_stat_counter=0)

        if self.DDP:
            self.group = dist.new_group([n for n in range(self.OPTIMIZATION.world_size)])

        if self.RUN.mixed_precision and not self.is_stylegan:
            self.scaler = torch.cuda.amp.GradScaler()

        if self.global_rank == 0:
            resume = False if self.RUN.freezeD > -1 else True
            wandb.init(project=self.RUN.project,
                       entity=self.RUN.entity,
                       name=self.run_name,
                       dir=self.RUN.save_dir,
                       resume=self.best_step > 0 and resume)

        self.start_time = datetime.now()

    def prepare_train_iter(self, epoch_counter):
        self.epoch_counter = epoch_counter
        if self.DDP:
            self.train_dataloader.sampler.set_epoch(self.epoch_counter)
        self.train_iter = iter(self.train_dataloader)

    def sample_data_basket(self):
        try:
            real_image_basket, real_label_basket = next(self.train_iter)
        except StopIteration:
            self.epoch_counter += 1
            if self.RUN.train and self.DDP:
                self.train_dataloader.sampler.set_epoch(self.epoch_counter)
            else:
                pass
            self.train_iter = iter(self.train_dataloader)
            real_image_basket, real_label_basket = next(self.train_iter)

        real_image_basket = torch.split(real_image_basket, self.OPTIMIZATION.batch_size)
        real_label_basket = torch.split(real_label_basket, self.OPTIMIZATION.batch_size)
        return real_image_basket, real_label_basket

    # -----------------------------------------------------------------------------
    # train Discriminator
    # -----------------------------------------------------------------------------
    def train_discriminator(self, current_step):
        batch_counter = 0
        # make GAN be trainable before starting training
        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)
        # toggle gradients of the generator and discriminator
        misc.toggle_grad(model=self.Gen, grad=False, num_freeze_layers=-1, is_stylegan=self.is_stylegan)
        misc.toggle_grad(model=self.Dis, grad=True, num_freeze_layers=self.RUN.freezeD, is_stylegan=self.is_stylegan)
        if self.MODEL.info_type in ["discrete", "both"]:
            misc.toggle_grad(getattr(misc.peel_model(self.Dis), self.MISC.info_params[0]), grad=False, num_freeze_layers=-1, is_stylegan=False)
        if self.MODEL.info_type in ["continuous", "both"]:
            misc.toggle_grad(getattr(misc.peel_model(self.Dis), self.MISC.info_params[1]), grad=False, num_freeze_layers=-1, is_stylegan=False)
            misc.toggle_grad(getattr(misc.peel_model(self.Dis), self.MISC.info_params[2]), grad=False, num_freeze_layers=-1, is_stylegan=False)
        if self.DDP*self.RUN.mixed_precision*self.RUN.synchronized_bn == 0: self.Gen.apply(misc.untrack_bn_statistics)
        # sample real images and labels from the true data distribution
        real_image_basket, real_label_basket = self.sample_data_basket()
        for step_index in range(self.OPTIMIZATION.d_updates_per_step):
            self.OPTIMIZATION.d_optimizer.zero_grad()
            for acml_index in range(self.OPTIMIZATION.acml_steps):
                with torch.cuda.amp.autocast() if self.RUN.mixed_precision and not self.is_stylegan else misc.dummy_context_mgr() as mpc:
                    # load real images and labels onto the GPU memory
                    real_images = real_image_basket[batch_counter].to(self.local_rank, non_blocking=True)
                    real_labels = real_label_basket[batch_counter].to(self.local_rank, non_blocking=True)
                    # sample fake images and labels from p(G(z), y)
                    fake_images, fake_labels, fake_images_eps, trsp_cost, ws, _, _ = sample.generate_images(
                        z_prior=self.MODEL.z_prior,
                        truncation_factor=-1.0,
                        batch_size=self.OPTIMIZATION.batch_size,
                        z_dim=self.MODEL.z_dim,
                        num_classes=self.DATA.num_classes,
                        y_sampler="totally_random",
                        radius=self.LOSS.radius,
                        generator=self.Gen,
                        discriminator=self.Dis,
                        is_train=True,
                        LOSS=self.LOSS,
                        RUN=self.RUN,
                        MODEL=self.MODEL,
                        device=self.local_rank,
                        generator_mapping=self.Gen_mapping,
                        generator_synthesis=self.Gen_synthesis,
                        is_stylegan=self.is_stylegan,
                        style_mixing_p=self.cfgs.STYLEGAN.style_mixing_p,
                        stylegan_update_emas=True,
                        cal_trsp_cost=True if self.LOSS.apply_lo else False)

                    # if LOSS.apply_r1_reg is True,
                    # let real images require gradient calculation to compute \derv_{x}Dis(x)
                    if self.LOSS.apply_r1_reg and not self.is_stylegan:
                        real_images.requires_grad_(True)

                    # blur images for stylegan3-r
                    if self.MODEL.backbone == "stylegan3" and self.STYLEGAN.stylegan3_cfg == "stylegan3-r" and self.blur_init_sigma != "N/A":
                        blur_sigma = max(1 - (self.effective_batch_size * current_step) / (self.blur_fade_kimg * 1e3), 0) * self.blur_init_sigma
                        blur_size = np.floor(blur_sigma * 3)
                        if blur_size > 0:
                            f = torch.arange(-blur_size, blur_size + 1, device=real_images.device).div(blur_sigma).square().neg().exp2()
                            real_images = upfirdn2d.filter2d(real_images, f / f.sum())
                            fake_images = upfirdn2d.filter2d(fake_images, f / f.sum())

                    # shuffle real and fake images (APA)
                    if self.AUG.apply_apa:
                        real_images = apa_aug.apply_apa_aug(real_images, fake_images.detach(), self.aa_p, self.local_rank)

                    # apply differentiable augmentations if "apply_diffaug" or "apply_ada" is True
                    real_images_ = self.AUG.series_augment(real_images)
                    fake_images_ = self.AUG.series_augment(fake_images)

                    # calculate adv_output, embed, proxy, and cls_output using the discriminator
                    real_dict = self.Dis(real_images_, real_labels)
                    fake_dict = self.Dis(fake_images_, fake_labels, adc_fake=self.adc_fake)

                    # accumulate discriminator output informations for logging
                    if self.AUG.apply_ada or self.AUG.apply_apa:
                        self.dis_sign_real += torch.tensor((real_dict["adv_output"].sign().sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                        device=self.local_rank)
                        self.dis_sign_fake += torch.tensor((fake_dict["adv_output"].sign().sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                        device=self.local_rank)
                        self.dis_logit_real += torch.tensor((real_dict["adv_output"].sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                            device=self.local_rank)
                        self.dis_logit_fake += torch.tensor((fake_dict["adv_output"].sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                            device=self.local_rank)

                    # calculate adversarial loss defined by "LOSS.adv_loss"
                    if self.LOSS.adv_loss == "MH":
                        dis_acml_loss = self.LOSS.d_loss(DDP=self.DDP, **real_dict)
                        dis_acml_loss += self.LOSS.d_loss(fake_dict["adv_output"], self.lossy, DDP=self.DDP)
                    else:
                        dis_acml_loss = self.LOSS.d_loss(real_dict["adv_output"], fake_dict["adv_output"], DDP=self.DDP)

                    # calculate class conditioning loss defined by "MODEL.d_cond_mtd"
                    if self.MODEL.d_cond_mtd in self.MISC.classifier_based_GAN:
                        real_cond_loss = self.cond_loss(**real_dict)
                        dis_acml_loss += self.LOSS.cond_lambda * real_cond_loss
                        if self.MODEL.aux_cls_type == "TAC":
                            tac_dis_loss = self.cond_loss_mi(**fake_dict)
                            dis_acml_loss += self.LOSS.tac_dis_lambda * tac_dis_loss
                        elif self.MODEL.aux_cls_type == "ADC":
                            fake_cond_loss = self.cond_loss(**fake_dict)
                            dis_acml_loss += self.LOSS.cond_lambda * fake_cond_loss
                        else:
                            pass
                    else:
                        real_cond_loss = "N/A"

                    # add transport cost for latent optimization training
                    if self.LOSS.apply_lo:
                        dis_acml_loss += self.LOSS.lo_lambda * trsp_cost

                    # if LOSS.apply_cr is True, force the adv. and cls. logits to be the same
                    if self.LOSS.apply_cr:
                        real_prl_images = self.AUG.parallel_augment(real_images)
                        real_prl_dict = self.Dis(real_prl_images, real_labels)
                        real_consist_loss = self.l2_loss(real_dict["adv_output"], real_prl_dict["adv_output"])
                        if self.MODEL.d_cond_mtd == "AC":
                            real_consist_loss += self.l2_loss(real_dict["cls_output"], real_prl_dict["cls_output"])
                        elif self.MODEL.d_cond_mtd in ["2C", "D2DCE"]:
                            real_consist_loss += self.l2_loss(real_dict["embed"], real_prl_dict["embed"])
                        else:
                            pass
                        dis_acml_loss += self.LOSS.cr_lambda * real_consist_loss

                    # if LOSS.apply_bcr is True, apply balanced consistency regularization proposed in ICRGAN
                    if self.LOSS.apply_bcr:
                        real_prl_images = self.AUG.parallel_augment(real_images)
                        fake_prl_images = self.AUG.parallel_augment(fake_images)
                        real_prl_dict = self.Dis(real_prl_images, real_labels)
                        fake_prl_dict = self.Dis(fake_prl_images, fake_labels, adc_fake=self.adc_fake)
                        real_bcr_loss = self.l2_loss(real_dict["adv_output"], real_prl_dict["adv_output"])
                        fake_bcr_loss = self.l2_loss(fake_dict["adv_output"], fake_prl_dict["adv_output"])
                        if self.MODEL.d_cond_mtd == "AC":
                            real_bcr_loss += self.l2_loss(real_dict["cls_output"], real_prl_dict["cls_output"])
                            fake_bcr_loss += self.l2_loss(fake_dict["cls_output"], fake_prl_dict["cls_output"])
                        elif self.MODEL.d_cond_mtd in ["2C", "D2DCE"]:
                            real_bcr_loss += self.l2_loss(real_dict["embed"], real_prl_dict["embed"])
                            fake_bcr_loss += self.l2_loss(fake_dict["embed"], fake_prl_dict["embed"])
                        else:
                            pass
                        dis_acml_loss += self.LOSS.real_lambda * real_bcr_loss + self.LOSS.fake_lambda * fake_bcr_loss

                    # if LOSS.apply_zcr is True, apply latent consistency regularization proposed in ICRGAN
                    if self.LOSS.apply_zcr:
                        fake_eps_dict = self.Dis(fake_images_eps, fake_labels, adc_fake=self.adc_fake)
                        fake_zcr_loss = self.l2_loss(fake_dict["adv_output"], fake_eps_dict["adv_output"])
                        if self.MODEL.d_cond_mtd == "AC":
                            fake_zcr_loss += self.l2_loss(fake_dict["cls_output"], fake_eps_dict["cls_output"])
                        elif self.MODEL.d_cond_mtd in ["2C", "D2DCE"]:
                            fake_zcr_loss += self.l2_loss(fake_dict["embed"], fake_eps_dict["embed"])
                        else:
                            pass
                        dis_acml_loss += self.LOSS.d_lambda * fake_zcr_loss

                    # apply gradient penalty regularization to train wasserstein GAN
                    if self.LOSS.apply_gp:
                        gp_loss = losses.cal_grad_penalty(real_images=real_images,
                                                          real_labels=real_labels,
                                                          fake_images=fake_images,
                                                          discriminator=self.Dis,
                                                          device=self.local_rank)
                        dis_acml_loss += self.LOSS.gp_lambda * gp_loss

                    # apply deep regret analysis regularization to train wasserstein GAN
                    if self.LOSS.apply_dra:
                        dra_loss = losses.cal_dra_penalty(real_images=real_images,
                                                          real_labels=real_labels,
                                                          discriminator=self.Dis,
                                                          device=self.local_rank)
                        dis_acml_loss += self.LOSS.dra_lambda * dra_loss

                    # apply max gradient penalty regularization to train Lipschitz GAN
                    if self.LOSS.apply_maxgp:
                        maxgp_loss = losses.cal_maxgrad_penalty(real_images=real_images,
                                                                real_labels=real_labels,
                                                                fake_images=fake_images,
                                                                discriminator=self.Dis,
                                                                device=self.local_rank)
                        dis_acml_loss += self.LOSS.maxgp_lambda * maxgp_loss

                    # apply LeCam reg. for data-efficient training if self.LOSS.apply_lecam is set to True
                    if self.LOSS.apply_lecam:
                        if self.DDP:
                            real_adv_output = torch.cat(losses.GatherLayer.apply(real_dict["adv_output"]), dim=0)
                            fake_adv_output = torch.cat(losses.GatherLayer.apply(fake_dict["adv_output"]), dim=0)
                        else:
                            real_adv_output, fake_adv_output = real_dict["adv_output"], fake_dict["adv_output"]
                        self.lecam_ema.update(torch.mean(real_adv_output).item(), "D_real", current_step)
                        self.lecam_ema.update(torch.mean(fake_adv_output).item(), "D_fake", current_step)
                        if current_step > self.LOSS.lecam_ema_start_iter:
                            lecam_loss = losses.lecam_reg(real_adv_output, fake_adv_output, self.lecam_ema)
                        else:
                            lecam_loss = torch.tensor(0., device=self.local_rank)
                        dis_acml_loss += self.LOSS.lecam_lambda*lecam_loss

                    # apply r1_reg inside of training loop
                    if self.LOSS.apply_r1_reg and not self.is_stylegan:
                        self.r1_penalty = losses.cal_r1_reg(adv_output=real_dict["adv_output"], images=real_images, device=self.local_rank)
                        dis_acml_loss += self.LOSS.r1_lambda*self.r1_penalty
                    elif self.LOSS.apply_r1_reg and self.LOSS.r1_place == "inside_loop" and \
                        (self.OPTIMIZATION.d_updates_per_step*current_step + step_index) % self.STYLEGAN.d_reg_interval == 0:
                        real_images.requires_grad_(True)
                        real_dict = self.Dis(self.AUG.series_augment(real_images), real_labels)
                        self.r1_penalty = losses.stylegan_cal_r1_reg(adv_output=real_dict["adv_output"],
                                                                     images=real_images)
                        dis_acml_loss += self.STYLEGAN.d_reg_interval*self.LOSS.r1_lambda*self.r1_penalty
                        if self.AUG.apply_ada or self.AUG.apply_apa:
                            self.dis_sign_real += torch.tensor((real_dict["adv_output"].sign().sum().item(),
                                                                self.OPTIMIZATION.batch_size),
                                                            device=self.local_rank)
                            self.dis_logit_real += torch.tensor((real_dict["adv_output"].sum().item(),
                                                                self.OPTIMIZATION.batch_size),
                                                                device=self.local_rank)

                    # adjust gradients for applying gradient accumluation trick
                    dis_acml_loss = dis_acml_loss / self.OPTIMIZATION.acml_steps
                    batch_counter += 1

                # accumulate gradients of the discriminator
                if self.RUN.mixed_precision and not self.is_stylegan:
                    self.scaler.scale(dis_acml_loss).backward()
                else:
                    dis_acml_loss.backward()

            # update the discriminator using the pre-defined optimizer
            if self.RUN.mixed_precision and not self.is_stylegan:
                self.scaler.step(self.OPTIMIZATION.d_optimizer)
                self.scaler.update()
            else:
                self.OPTIMIZATION.d_optimizer.step()

            # apply r1_reg outside of training loop
            if self.LOSS.apply_r1_reg and self.LOSS.r1_place == "outside_loop" and \
             (self.OPTIMIZATION.d_updates_per_step*current_step + step_index) % self.STYLEGAN.d_reg_interval == 0:
                self.OPTIMIZATION.d_optimizer.zero_grad()
                for acml_index in range(self.OPTIMIZATION.acml_steps):
                    real_images = real_image_basket[batch_counter - acml_index - 1].to(self.local_rank, non_blocking=True)
                    real_labels = real_label_basket[batch_counter - acml_index - 1].to(self.local_rank, non_blocking=True)
                    # blur images for stylegan3-r
                    if self.MODEL.backbone == "stylegan3" and self.STYLEGAN.stylegan3_cfg == "stylegan3-r" and self.blur_init_sigma != "N/A":
                        blur_sigma = max(1 - (self.effective_batch_size * current_step) / (self.blur_fade_kimg * 1e3), 0) * self.blur_init_sigma
                        blur_size = np.floor(blur_sigma * 3)
                        if blur_size > 0:
                            f = torch.arange(-blur_size, blur_size + 1, device=real_images.device).div(blur_sigma).square().neg().exp2()
                            real_images = upfirdn2d.filter2d(real_images, f / f.sum())
                    if self.AUG.apply_apa:
                        real_images = apa_aug.apply_apa_aug(real_images, fake_images.detach(), self.aa_p, self.local_rank)
                    real_images.requires_grad_(True)
                    real_dict = self.Dis(self.AUG.series_augment(real_images), real_labels)
                    self.r1_penalty = losses.stylegan_cal_r1_reg(adv_output=real_dict["adv_output"], images=real_images) + \
                        misc.enable_allreduce(real_dict)
                    self.r1_penalty *= self.STYLEGAN.d_reg_interval*self.LOSS.r1_lambda/self.OPTIMIZATION.acml_steps
                    self.r1_penalty.backward()

                    if self.AUG.apply_ada or self.AUG.apply_apa:
                        self.dis_sign_real += torch.tensor((real_dict["adv_output"].sign().sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                        device=self.local_rank)
                        self.dis_logit_real += torch.tensor((real_dict["adv_output"].sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                            device=self.local_rank)
                self.OPTIMIZATION.d_optimizer.step()

            # apply ada heuristics
            if (self.AUG.apply_ada or self.AUG.apply_apa) and self.aa_target is not None and current_step % self.aa_interval == 0:
                if self.DDP: dist.all_reduce(self.dis_sign_real, op=dist.ReduceOp.SUM, group=self.group)
                heuristic = (self.dis_sign_real[0] / self.dis_sign_real[1]).item()
                adjust = np.sign(heuristic - self.aa_target) * (self.dis_sign_real[1].item()) / (self.aa_kimg * 1000)
                self.aa_p = min(torch.as_tensor(1.), max(self.aa_p + adjust, torch.as_tensor(0.)))
                if self.AUG.apply_ada: self.AUG.series_augment.p.copy_(torch.as_tensor(self.aa_p))
                self.dis_sign_real_log.copy_(self.dis_sign_real), self.dis_sign_fake_log.copy_(self.dis_sign_fake)
                self.dis_logit_real_log.copy_(self.dis_logit_real), self.dis_logit_fake_log.copy_(self.dis_logit_fake)
                self.dis_sign_real.mul_(0), self.dis_sign_fake.mul_(0)
                self.dis_logit_real.mul_(0), self.dis_logit_fake.mul_(0)

            # clip weights to restrict the discriminator to satisfy 1-Lipschitz constraint
            if self.LOSS.apply_wc:
                for p in self.Dis.parameters():
                    p.data.clamp_(-self.LOSS.wc_bound, self.LOSS.wc_bound)

        # empty cache to discard used memory
        if self.RUN.empty_cache:
            torch.cuda.empty_cache()
        return real_cond_loss, dis_acml_loss

    # -----------------------------------------------------------------------------
    # train Generator
    # -----------------------------------------------------------------------------
    def train_generator(self, current_step):
        # make GAN be trainable before starting training
        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)
        # toggle gradients of the generator and discriminator
        misc.toggle_grad(model=self.Dis, grad=False, num_freeze_layers=-1, is_stylegan=self.is_stylegan)
        misc.toggle_grad(model=self.Gen, grad=True, num_freeze_layers=-1, is_stylegan=self.is_stylegan)
        if self.MODEL.info_type in ["discrete", "both"]:
            misc.toggle_grad(getattr(misc.peel_model(self.Dis), self.MISC.info_params[0]), grad=True, num_freeze_layers=-1, is_stylegan=False)
        if self.MODEL.info_type in ["continuous", "both"]:
            misc.toggle_grad(getattr(misc.peel_model(self.Dis), self.MISC.info_params[1]), grad=True, num_freeze_layers=-1, is_stylegan=False)
            misc.toggle_grad(getattr(misc.peel_model(self.Dis), self.MISC.info_params[2]), grad=True, num_freeze_layers=-1, is_stylegan=False)
        self.Gen.apply(misc.track_bn_statistics)
        for step_index in range(self.OPTIMIZATION.g_updates_per_step):
            self.OPTIMIZATION.g_optimizer.zero_grad()
            for acml_step in range(self.OPTIMIZATION.acml_steps):
                with torch.cuda.amp.autocast() if self.RUN.mixed_precision and not self.is_stylegan else misc.dummy_context_mgr() as mpc:
                    # sample fake images and labels from p(G(z), y)
                    fake_images, fake_labels, fake_images_eps, trsp_cost, ws, info_discrete_c, info_conti_c = sample.generate_images(
                        z_prior=self.MODEL.z_prior,
                        truncation_factor=-1.0,
                        batch_size=self.OPTIMIZATION.batch_size,
                        z_dim=self.MODEL.z_dim,
                        num_classes=self.DATA.num_classes,
                        y_sampler="totally_random",
                        radius=self.LOSS.radius,
                        generator=self.Gen,
                        discriminator=self.Dis,
                        is_train=True,
                        LOSS=self.LOSS,
                        RUN=self.RUN,
                        MODEL=self.MODEL,
                        device=self.local_rank,
                        generator_mapping=self.Gen_mapping,
                        generator_synthesis=self.Gen_synthesis,
                        is_stylegan=self.is_stylegan,
                        style_mixing_p=self.cfgs.STYLEGAN.style_mixing_p,
                        stylegan_update_emas=False,
                        cal_trsp_cost=True if self.LOSS.apply_lo else False)

                    # blur images for stylegan3-r
                    if self.MODEL.backbone == "stylegan3" and self.STYLEGAN.stylegan3_cfg == "stylegan3-r" and self.blur_init_sigma != "N/A":
                        blur_sigma = max(1 - (self.effective_batch_size * current_step) / (self.blur_fade_kimg * 1e3), 0) * self.blur_init_sigma
                        blur_size = np.floor(blur_sigma * 3)
                        if blur_size > 0:
                            f = torch.arange(-blur_size, blur_size + 1, device=fake_images.device).div(blur_sigma).square().neg().exp2()
                            fake_images = upfirdn2d.filter2d(fake_images, f / f.sum())

                    # apply differentiable augmentations if "apply_diffaug" is True
                    fake_images_ = self.AUG.series_augment(fake_images)

                    # calculate adv_output, embed, proxy, and cls_output using the discriminator
                    fake_dict = self.Dis(fake_images_, fake_labels)

                    # accumulate discriminator output informations for logging
                    if self.AUG.apply_ada or self.AUG.apply_apa:
                        self.dis_sign_fake += torch.tensor((fake_dict["adv_output"].sign().sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                        device=self.local_rank)
                        self.dis_logit_fake += torch.tensor((fake_dict["adv_output"].sum().item(),
                                                            self.OPTIMIZATION.batch_size),
                                                            device=self.local_rank)

                    # apply top k sampling for discarding bottom 1-k samples which are 'in-between modes'
                    if self.LOSS.apply_topk:
                        fake_dict["adv_output"] = torch.topk(fake_dict["adv_output"], int(self.topk)).values

                    # calculate adversarial loss defined by "LOSS.adv_loss"
                    if self.LOSS.adv_loss == "MH":
                        gen_acml_loss = self.LOSS.mh_lambda * self.LOSS.g_loss(DDP=self.DDP, **fake_dict, )
                    else:
                        gen_acml_loss = self.LOSS.g_loss(fake_dict["adv_output"], DDP=self.DDP)

                    # calculate class conditioning loss defined by "MODEL.d_cond_mtd"
                    if self.MODEL.d_cond_mtd in self.MISC.classifier_based_GAN:
                        fake_cond_loss = self.cond_loss(**fake_dict)
                        gen_acml_loss += self.LOSS.cond_lambda * fake_cond_loss
                        if self.MODEL.aux_cls_type == "TAC":
                            tac_gen_loss = -self.cond_loss_mi(**fake_dict)
                            gen_acml_loss += self.LOSS.tac_gen_lambda * tac_gen_loss
                        elif self.MODEL.aux_cls_type == "ADC":
                            adc_fake_dict = self.Dis(fake_images_, fake_labels, adc_fake=self.adc_fake)
                            adc_fake_cond_loss = -self.cond_loss(**adc_fake_dict)
                            gen_acml_loss += self.LOSS.cond_lambda * adc_fake_cond_loss
                        pass

                    # apply feature matching regularization to stabilize adversarial dynamics
                    if self.LOSS.apply_fm:
                        real_image_basket, real_label_basket = self.sample_data_basket()
                        real_images = real_image_basket[0].to(self.local_rank, non_blocking=True)
                        real_labels = real_label_basket[0].to(self.local_rank, non_blocking=True)
                        real_images_ = self.AUG.series_augment(real_images)
                        real_dict = self.Dis(real_images_, real_labels)

                        mean_match_loss = self.fm_loss(real_dict["h"].detach(), fake_dict["h"])
                        gen_acml_loss += self.LOSS.fm_lambda * mean_match_loss

                    # add transport cost for latent optimization training
                    if self.LOSS.apply_lo:
                        gen_acml_loss += self.LOSS.lo_lambda * trsp_cost

                    # apply latent consistency regularization for generating diverse images
                    if self.LOSS.apply_zcr:
                        fake_zcr_loss = -1 * self.l2_loss(fake_images, fake_images_eps)
                        gen_acml_loss += self.LOSS.g_lambda * fake_zcr_loss

                    # compute infomation loss for InfoGAN
                    if self.MODEL.info_type in ["discrete", "both"]:
                        dim = self.MODEL.info_dim_discrete_c
                        self.info_discrete_loss = 0.0
                        for info_c in range(self.MODEL.info_num_discrete_c):
                            self.info_discrete_loss += self.ce_loss(
                                fake_dict["info_discrete_c_logits"][:, info_c*dim: dim*(info_c+1)],
                                info_discrete_c[:, info_c: info_c+1].squeeze())
                        gen_acml_loss += self.LOSS.infoGAN_loss_discrete_lambda*self.info_discrete_loss + misc.enable_allreduce(fake_dict)
                    if self.MODEL.info_type in ["continuous", "both"]:
                        self.info_conti_loss = losses.normal_nll_loss(info_conti_c, fake_dict["info_conti_mu"], fake_dict["info_conti_var"])
                        gen_acml_loss += self.LOSS.infoGAN_loss_conti_lambda*self.info_conti_loss + misc.enable_allreduce(fake_dict)

                    # adjust gradients for applying gradient accumluation trick
                    gen_acml_loss = gen_acml_loss / self.OPTIMIZATION.acml_steps

                # accumulate gradients of the generator
                if self.RUN.mixed_precision and not self.is_stylegan:
                    self.scaler.scale(gen_acml_loss).backward()
                else:
                    gen_acml_loss.backward()

            # update the generator using the pre-defined optimizer
            if self.RUN.mixed_precision and not self.is_stylegan:
                self.scaler.step(self.OPTIMIZATION.g_optimizer)
                self.scaler.update()
            else:
                self.OPTIMIZATION.g_optimizer.step()

            # apply path length regularization
            if self.STYLEGAN.apply_pl_reg and (self.OPTIMIZATION.g_updates_per_step*current_step + step_index) % self.STYLEGAN.g_reg_interval == 0:
                self.OPTIMIZATION.g_optimizer.zero_grad()
                for acml_index in range(self.OPTIMIZATION.acml_steps):
                    fake_images, fake_labels, fake_images_eps, trsp_cost, ws, _, _ = sample.generate_images(
                        z_prior=self.MODEL.z_prior,
                        truncation_factor=-1.0,
                        batch_size=self.OPTIMIZATION.batch_size // 2,
                        z_dim=self.MODEL.z_dim,
                        num_classes=self.DATA.num_classes,
                        y_sampler="totally_random",
                        radius=self.LOSS.radius,
                        generator=self.Gen,
                        discriminator=self.Dis,
                        is_train=True,
                        LOSS=self.LOSS,
                        RUN=self.RUN,
                        MODEL=self.MODEL,
                        device=self.local_rank,
                        generator_mapping=self.Gen_mapping,
                        generator_synthesis=self.Gen_synthesis,
                        is_stylegan=self.is_stylegan,
                        style_mixing_p=self.cfgs.STYLEGAN.style_mixing_p,
                        stylegan_update_emas=False,
                        cal_trsp_cost=True if self.LOSS.apply_lo else False)

                    # blur images for stylegan3-r
                    if self.MODEL.backbone == "stylegan3" and self.STYLEGAN.stylegan3_cfg == "stylegan3-r" and self.blur_init_sigma != "N/A":
                        blur_sigma = max(1 - (self.effective_batch_size * current_step) / (self.blur_fade_kimg * 1e3), 0) * self.blur_init_sigma
                        blur_size = np.floor(blur_sigma * 3)
                        if blur_size > 0:
                            f = torch.arange(-blur_size, blur_size + 1, device=fake_images.device).div(blur_sigma).square().neg().exp2()
                            fake_images = upfirdn2d.filter2d(fake_images, f / f.sum())
                    self.pl_reg_loss = self.pl_reg.cal_pl_reg(fake_images=fake_images, ws=ws) + fake_images[:,0,0,0].mean()*0
                    self.pl_reg_loss *= self.STYLEGAN.g_reg_interval/self.OPTIMIZATION.acml_steps
                    self.pl_reg_loss.backward()
                self.OPTIMIZATION.g_optimizer.step()

            # if ema is True: update parameters of the Gen_ema in adaptive way
            if self.MODEL.apply_g_ema:
                self.ema.update(current_step)

        # empty cache to discard used memory
        if self.RUN.empty_cache:
            torch.cuda.empty_cache()
        return gen_acml_loss

    # -----------------------------------------------------------------------------
    # log training statistics
    # -----------------------------------------------------------------------------
    def log_train_statistics(self, current_step, real_cond_loss, gen_acml_loss, dis_acml_loss):
        self.wandb_step = current_step + 1
        if self.MODEL.d_cond_mtd in self.MISC.classifier_based_GAN:
            cls_loss = real_cond_loss.item()
        else:
            cls_loss = "N/A"

        log_message = LOG_FORMAT.format(
            step=current_step + 1,
            progress=(current_step + 1) / self.OPTIMIZATION.total_steps,
            elapsed=misc.elapsed_time(self.start_time),
            gen_loss=gen_acml_loss.item(),
            dis_loss=dis_acml_loss.item(),
            cls_loss=cls_loss,
            topk=int(self.topk) if self.LOSS.apply_topk else "N/A",
            aa_p=self.aa_p if self.AUG.apply_ada or self.AUG.apply_apa else "N/A",
        )
        self.logger.info(log_message)

        # save loss values in wandb event file and .npz format
        loss_dict = {
            "gen_loss": gen_acml_loss.item(),
            "dis_loss": dis_acml_loss.item(),
            "cls_loss": 0.0 if cls_loss == "N/A" else cls_loss,
        }

        wandb.log(loss_dict, step=self.wandb_step)

        save_dict = misc.accm_values_convert_dict(list_dict=self.loss_list_dict,
                                                  value_dict=loss_dict,
                                                  step=current_step + 1,
                                                  interval=self.RUN.print_freq)
        misc.save_dict_npy(directory=join(self.RUN.save_dir, "statistics", self.run_name),
                           name="losses",
                           dictionary=save_dict)

        if self.AUG.apply_ada or self.AUG.apply_apa:
            dis_output_dict = {
                        "dis_sign_real": (self.dis_sign_real_log[0]/self.dis_sign_real_log[1]).item(),
                        "dis_sign_fake": (self.dis_sign_fake_log[0]/self.dis_sign_fake_log[1]).item(),
                        "dis_logit_real": (self.dis_logit_real_log[0]/self.dis_logit_real_log[1]).item(),
                        "dis_logit_fake": (self.dis_logit_fake_log[0]/self.dis_logit_fake_log[1]).item(),
                    }
            wandb.log(dis_output_dict, step=self.wandb_step)
            wandb.log({"aa_p": self.aa_p.item()}, step=self.wandb_step)

        infoGAN_dict = {}
        if self.MODEL.info_type in ["discrete", "both"]:
            infoGAN_dict["info_discrete_loss"] = self.info_discrete_loss.item()
        if self.MODEL.info_type in ["continuous", "both"]:
            infoGAN_dict["info_conti_loss"] = self.info_conti_loss.item()
            wandb.log(infoGAN_dict, step=self.wandb_step)

        if self.LOSS.apply_r1_reg:
            wandb.log({"r1_reg_loss": self.r1_penalty.item()}, step=self.wandb_step)

        if self.STYLEGAN.apply_pl_reg:
            wandb.log({"pl_reg_loss": self.pl_reg_loss.item()}, step=self.wandb_step)

        # calculate the spectral norms of all weights in the generator for monitoring purpose
        if self.MODEL.apply_g_sn:
            gen_sigmas = misc.calculate_all_sn(self.Gen, prefix="Gen")
            wandb.log(gen_sigmas, step=self.wandb_step)

        # calculate the spectral norms of all weights in the discriminator for monitoring purpose
        if self.MODEL.apply_d_sn:
            dis_sigmas = misc.calculate_all_sn(self.Dis, prefix="Dis")
            wandb.log(dis_sigmas, step=self.wandb_step)

    # -----------------------------------------------------------------------------
    # visualize fake images for monitoring purpose.
    # -----------------------------------------------------------------------------
    def visualize_fake_images(self, num_cols, current_step):
        if self.global_rank == 0:
            self.logger.info("Visualize (num_rows x 8) fake image canvans.")
        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            fake_images, fake_labels, _, _, _, _, _ = sample.generate_images(z_prior=self.MODEL.z_prior,
                                                                       truncation_factor=self.RUN.truncation_factor,
                                                                       batch_size=self.OPTIMIZATION.batch_size,
                                                                       z_dim=self.MODEL.z_dim,
                                                                       num_classes=self.DATA.num_classes,
                                                                       y_sampler=self.sampler,
                                                                       radius="N/A",
                                                                       generator=generator,
                                                                       discriminator=self.Dis,
                                                                       is_train=False,
                                                                       LOSS=self.LOSS,
                                                                       RUN=self.RUN,
                                                                       MODEL=self.MODEL,
                                                                       device=self.local_rank,
                                                                       is_stylegan=self.is_stylegan,
                                                                       generator_mapping=generator_mapping,
                                                                       generator_synthesis=generator_synthesis,
                                                                       style_mixing_p=0.0,
                                                                       stylegan_update_emas=False,
                                                                       cal_trsp_cost=False)

        misc.plot_img_canvas(images=fake_images.detach().cpu(),
                             save_path=join(self.RUN.save_dir,
                             "figures/{run_name}/generated_canvas_{step}.png".format(run_name=self.run_name, step=current_step)),
                             num_cols=num_cols,
                             logger=self.logger,
                             logging=self.global_rank == 0 and self.logger)

        if self.RUN.train and os.environ.get("WANDB_MODE", "").lower() not in ("disabled", "offline") \
                and os.environ.get("WANDB_DISABLED", "").lower() not in ("1", "true", "yes"):
            wandb.log({"generated_images": wandb.Image(fake_images)}, step=self.wandb_step)

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # evaluate GAN using IS, FID, and Precision and recall.
    # -----------------------------------------------------------------------------
    def evaluate(self, step, metrics, writing=True, training=False):
        if self.global_rank == 0:
            self.logger.info("Start Evaluation ({step} Step): {run_name}".format(step=step, run_name=self.run_name))
        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        is_best, num_splits, nearest_k = False, 1, 5
        is_acc = True if "ImageNet" in self.DATA.name and "Tiny" not in self.DATA.name else False
        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()
            metric_dict = {}

            fake_feats, fake_probs, fake_labels = features.generate_images_and_stack_features(
                                                                   generator=generator,
                                                                   discriminator=self.Dis,
                                                                   eval_model=self.eval_model,
                                                                   num_generate=self.num_eval[self.RUN.ref_dataset],
                                                                   y_sampler="totally_random",
                                                                   batch_size=self.OPTIMIZATION.batch_size,
                                                                   z_prior=self.MODEL.z_prior,
                                                                   truncation_factor=self.RUN.truncation_factor,
                                                                   z_dim=self.MODEL.z_dim,
                                                                   num_classes=self.DATA.num_classes,
                                                                   LOSS=self.LOSS,
                                                                   RUN=self.RUN,
                                                                   MODEL=self.MODEL,
                                                                   is_stylegan=self.is_stylegan,
                                                                   generator_mapping=generator_mapping,
                                                                   generator_synthesis=generator_synthesis,
                                                                   quantize=True,
                                                                   world_size=self.OPTIMIZATION.world_size,
                                                                   DDP=self.DDP,
                                                                   device=self.local_rank,
                                                                   logger=self.logger,
                                                                   disable_tqdm=self.global_rank != 0)

            if ("fid" in metrics or "prdc" in metrics) and self.global_rank == 0:
                self.logger.info("{num_images} real images is used for evaluation.".format(num_images=len(self.eval_dataloader.dataset)))

            if "is" in metrics:
                kl_score, kl_std, top1, top5 = ins.eval_features(probs=fake_probs,
                                                                 labels=fake_labels,
                                                                 data_loader=self.eval_dataloader,
                                                                 num_features=self.num_eval[self.RUN.ref_dataset],
                                                                 split=num_splits,
                                                                 is_acc=is_acc,
                                                                 is_torch_backbone=True if "torch" in self.RUN.eval_backbone else False)
                if self.global_rank == 0:
                    self.logger.info("Inception score (Step: {step}, {num} generated images): {IS}".format(
                        step=step, num=str(self.num_eval[self.RUN.ref_dataset]), IS=kl_score))
                    if is_acc:
                        self.logger.info("{eval_model} Top1 acc: (Step: {step}, {num} generated images): {Top1}".format(
                            eval_model=self.RUN.eval_backbone, step=step, num=str(self.num_eval[self.RUN.ref_dataset]), Top1=top1))
                        self.logger.info("{eval_model} Top5 acc: (Step: {step}, {num} generated images): {Top5}".format(
                            eval_model=self.RUN.eval_backbone, step=step, num=str(self.num_eval[self.RUN.ref_dataset]), Top5=top5))
                    metric_dict.update({"IS": kl_score, "Top1_acc": top1, "Top5_acc": top5})
                    if writing:
                        wandb.log({"IS score": kl_score}, step=self.wandb_step)
                        if is_acc:
                            wandb.log({"{eval_model} Top1 acc".format(eval_model=self.RUN.eval_backbone): top1}, step=self.wandb_step)
                            wandb.log({"{eval_model} Top5 acc".format(eval_model=self.RUN.eval_backbone): top5}, step=self.wandb_step)

            if "fid" in metrics:
                fid_score, m1, c1 = fid.calculate_fid(data_loader=self.eval_dataloader,
                                                      eval_model=self.eval_model,
                                                      num_generate=self.num_eval[self.RUN.ref_dataset],
                                                      cfgs=self.cfgs,
                                                      pre_cal_mean=self.mu,
                                                      pre_cal_std=self.sigma,
                                                      fake_feats=fake_feats,
                                                      disable_tqdm=self.global_rank != 0)
                if self.global_rank == 0:
                    self.logger.info("FID score (Step: {step}, Using {type} moments): {FID}".format(
                        step=step, type=self.RUN.ref_dataset, FID=fid_score))
                    if self.best_fid is None or fid_score <= self.best_fid:
                        self.best_fid, self.best_step, is_best = fid_score, step, True
                    metric_dict.update({"FID": fid_score})
                    if writing:
                        wandb.log({"FID score": fid_score}, step=self.wandb_step)
                    if training:
                        self.logger.info("Best FID score (Step: {step}, Using {type} moments): {FID}".format(
                            step=self.best_step, type=self.RUN.ref_dataset, FID=self.best_fid))

            if "prdc" in metrics:
                prc, rec, dns, cvg = prdc.calculate_pr_dc(real_feats=self.real_feats,
                                                          fake_feats=fake_feats,
                                                          data_loader=self.eval_dataloader,
                                                          eval_model=self.eval_model,
                                                          num_generate=self.num_eval[self.RUN.ref_dataset],
                                                          cfgs=self.cfgs,
                                                          quantize=True,
                                                          nearest_k=nearest_k,
                                                          world_size=self.OPTIMIZATION.world_size,
                                                          DDP=self.DDP,
                                                          disable_tqdm=True)
                if self.global_rank == 0:
                    self.logger.info("Improved Precision (Step: {step}, Using {type} images): {prc}".format(
                        step=step, type=self.RUN.ref_dataset, prc=prc))
                    self.logger.info("Improved Recall (Step: {step}, Using {type} images): {rec}".format(
                        step=step, type=self.RUN.ref_dataset, rec=rec))
                    self.logger.info("Density (Step: {step}, Using {type} images): {dns}".format(
                        step=step, type=self.RUN.ref_dataset, dns=dns))
                    self.logger.info("Coverage (Step: {step}, Using {type} images): {cvg}".format(
                        step=step, type=self.RUN.ref_dataset, cvg=cvg))
                    metric_dict.update({"Improved_Precision": prc, "Improved_Recall": rec, "Density": dns, "Coverage": cvg})
                    if writing:
                        wandb.log({"Improved Precision": prc}, step=self.wandb_step)
                        wandb.log({"Improved Recall": rec}, step=self.wandb_step)
                        wandb.log({"Density": dns}, step=self.wandb_step)
                        wandb.log({"Coverage": cvg}, step=self.wandb_step)

            if self.global_rank == 0:
                if training:
                    save_dict = misc.accm_values_convert_dict(list_dict=self.metric_dict_during_train,
                                                              value_dict=metric_dict,
                                                              step=step,
                                                              interval=self.RUN.save_freq)
                else:
                    save_dict = misc.accm_values_convert_dict(list_dict=self.metric_dict_during_final_eval,
                                                              value_dict=metric_dict,
                                                              step=None,
                                                              interval=None)

                    misc.save_dict_npy(directory=join(self.RUN.save_dir, "statistics", self.run_name, "train" if training else "eval"),
                                       name="metrics",
                                       dictionary=save_dict)

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)
        return is_best

    # -----------------------------------------------------------------------------
    # save the trained generator, generator_ema, and discriminator.
    # -----------------------------------------------------------------------------
    def save(self, step, is_best):
        when = "best" if is_best is True else "current"
        misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
        Gen, Gen_ema, Dis = misc.peel_models(self.Gen, self.Gen_ema, self.Dis)

        g_states = {"state_dict": Gen.state_dict(), "optimizer": self.OPTIMIZATION.g_optimizer.state_dict()}

        d_states = {
            "state_dict": Dis.state_dict(),
            "optimizer": self.OPTIMIZATION.d_optimizer.state_dict(),
            "seed": self.RUN.seed,
            "run_name": self.run_name,
            "step": step,
            "epoch": self.epoch_counter,
            "topk": self.topk,
            "aa_p": self.aa_p,
            "best_step": self.best_step,
            "best_fid": self.best_fid,
            "best_fid_ckpt": self.RUN.ckpt_dir,
            "lecam_emas": self.lecam_ema.__dict__,
        }

        if self.Gen_ema is not None:
            g_ema_states = {"state_dict": Gen_ema.state_dict()}

        misc.save_model(model="G", when=when, step=step, ckpt_dir=self.RUN.ckpt_dir, states=g_states)
        misc.save_model(model="D", when=when, step=step, ckpt_dir=self.RUN.ckpt_dir, states=d_states)
        if self.Gen_ema is not None:
            misc.save_model(model="G_ema", when=when, step=step, ckpt_dir=self.RUN.ckpt_dir, states=g_ema_states)

        if when == "best":
            misc.save_model(model="G", when="current", step=step, ckpt_dir=self.RUN.ckpt_dir, states=g_states)
            misc.save_model(model="D", when="current", step=step, ckpt_dir=self.RUN.ckpt_dir, states=d_states)
            if self.Gen_ema is not None:
                misc.save_model(model="G_ema",
                                when="current",
                                step=step,
                                ckpt_dir=self.RUN.ckpt_dir,
                                states=g_ema_states)

        if self.global_rank == 0 and self.logger:
            self.logger.info("Save model to {}".format(self.RUN.ckpt_dir))

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # save real images to measure metrics for evaluation.
    # -----------------------------------------------------------------------------
    def save_real_images(self):
        if self.global_rank == 0:
            self.logger.info("save {num_images} real images in png format.".format(
                num_images=len(self.eval_dataloader.dataset)))

        misc.save_images_png(data_loader=self.eval_dataloader,
                             generator="N/A",
                             discriminator="N/A",
                             is_generate=False,
                             num_images=len(self.eval_dataloader.dataset),
                             y_sampler="N/A",
                             batch_size=self.OPTIMIZATION.batch_size,
                             z_prior="N/A",
                             truncation_factor="N/A",
                             z_dim="N/A",
                             num_classes=self.DATA.num_classes,
                             LOSS=self.LOSS,
                             OPTIMIZATION=self.OPTIMIZATION,
                             RUN=self.RUN,
                             MODEL=self.MODEL,
                             is_stylegan=False,
                             generator_mapping="N/A",
                             generator_synthesis="N/A",
                             directory=join(self.RUN.save_dir, "samples", self.run_name),
                             device=self.local_rank)

    # -----------------------------------------------------------------------------
    # save fake images to measure metrics for evaluation.
    # -----------------------------------------------------------------------------
    def save_fake_images(self, num_images):
        if self.global_rank == 0:
            self.logger.info("save {num_images} generated images in png format.".format(num_images=self.num_eval[self.RUN.ref_dataset]))
        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            misc.save_images_png(data_loader=self.eval_dataloader,
                                 generator=generator,
                                 discriminator=self.Dis,
                                 is_generate=True,
                                 num_images=num_images,
                                 y_sampler="totally_random",
                                 batch_size=self.OPTIMIZATION.batch_size,
                                 z_prior=self.MODEL.z_prior,
                                 truncation_factor=self.RUN.truncation_factor,
                                 z_dim=self.MODEL.z_dim,
                                 num_classes=self.DATA.num_classes,
                                 LOSS=self.LOSS,
                                 OPTIMIZATION=self.OPTIMIZATION,
                                 RUN=self.RUN,
                                 MODEL=self.MODEL,
                                 is_stylegan=self.is_stylegan,
                                 generator_mapping=generator_mapping,
                                 generator_synthesis=generator_synthesis,
                                 directory=join(self.RUN.save_dir, "samples", self.run_name),
                                 device=self.local_rank)

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # run k-nearest neighbor analysis to identify whether GAN memorizes the training images or not.
    # -----------------------------------------------------------------------------
    def run_k_nearest_neighbor(self, dataset, num_rows, num_cols):
        if self.global_rank == 0:
            self.logger.info("Run K-nearest neighbor analysis using fake and {ref} dataset.".format(ref=self.RUN.ref_dataset))
        if self.gen_ctlr.standing_statistics: self.gen_ctlr.std_stat_counter += 1

        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            res, mean, std = 224, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
            resizer = resize.build_resizer(resizer=self.RUN.post_resizer,
                                           backbone="ResNet50_torch",
                                           size=res)
            totensor = transforms.ToTensor()
            mean = torch.Tensor(mean).view(1, 3, 1, 1).to("cuda")
            std = torch.Tensor(std).view(1, 3, 1, 1).to("cuda")

            # torch.hub.load("pytorch/vision:v0.6.0", ...) clones the GitHub
            # repo at runtime (to read its hubconf.py), which times out on
            # firewalled compute nodes. Build the architecture directly via
            # torchvision and load the cached weights file -- this is the same
            # content torch.hub would have downloaded.
            _resnet50_basename = "resnet50-19c8e357.pth"
            _resnet50_candidates = []
            _torch_home_env = os.environ.get("TORCH_HOME")
            if _torch_home_env:
                _resnet50_candidates += [
                    os.path.join(_torch_home_env, "hub", "checkpoints", _resnet50_basename),
                    os.path.join(_torch_home_env, "checkpoints", _resnet50_basename),
                ]
            _resnet50_candidates += [
                os.path.join(torch.hub.get_dir(), "checkpoints", _resnet50_basename),
                os.path.expanduser(os.path.join("~", ".cache", "torch", "hub",
                                                "checkpoints", _resnet50_basename)),
            ]
            _resnet50_ckpt = next((p for p in _resnet50_candidates if os.path.isfile(p)), None)
            if _resnet50_ckpt is not None:
                print("[knn] using cached resnet50 weights at:", _resnet50_ckpt, flush=True)
                resnet50_model = torchvision.models.resnet50(weights=None)
                resnet50_model.load_state_dict(
                    torch.load(_resnet50_ckpt, map_location="cpu", weights_only=False)
                )
            else:
                print("[knn] no cached resnet50 weights found in:",
                      _resnet50_candidates, "-- falling back to torch.hub.load",
                      flush=True)
                resnet50_model = torch.hub.load("pytorch/vision:v0.6.0",
                                                "resnet50", pretrained=True)
            resnet50_conv = nn.Sequential(*list(resnet50_model.children())[:-1]).to(self.local_rank)
            if self.OPTIMIZATION.world_size > 1:
                resnet50_conv = DataParallel(resnet50_conv, output_device=self.local_rank)
            resnet50_conv.eval()

            for c in tqdm(range(self.DATA.num_classes)):
                fake_images, fake_labels, _, _, _, _, _ = sample.generate_images(z_prior=self.MODEL.z_prior,
                                                                        truncation_factor=self.RUN.truncation_factor,
                                                                        batch_size=self.OPTIMIZATION.batch_size,
                                                                        z_dim=self.MODEL.z_dim,
                                                                        num_classes=self.DATA.num_classes,
                                                                        y_sampler=c,
                                                                        radius="N/A",
                                                                        generator=generator,
                                                                        discriminator=self.Dis,
                                                                        is_train=False,
                                                                        LOSS=self.LOSS,
                                                                        RUN=self.RUN,
                                                                        MODEL=self.MODEL,
                                                                        device=self.local_rank,
                                                                        is_stylegan=self.is_stylegan,
                                                                        generator_mapping=generator_mapping,
                                                                        generator_synthesis=generator_synthesis,
                                                                        style_mixing_p=0.0,
                                                                        stylegan_update_emas=False,
                                                                        cal_trsp_cost=False)
                fake_anchor = torch.unsqueeze(fake_images[0], dim=0)
                fake_anchor = ops.quantize_images(fake_anchor)
                fake_anchor = ops.resize_images(fake_anchor, resizer, totensor, mean, std, self.local_rank)
                fake_anchor_embed = torch.squeeze(resnet50_conv(fake_anchor))

                num_samples, target_sampler = sample.make_target_cls_sampler(dataset=dataset, target_class=c)
                batch_size = self.OPTIMIZATION.batch_size if num_samples >= self.OPTIMIZATION.batch_size else num_samples
                c_dataloader = torch.utils.data.DataLoader(dataset=dataset,
                                                           batch_size=batch_size,
                                                           shuffle=False,
                                                           sampler=target_sampler,
                                                           num_workers=self.RUN.num_workers,
                                                           pin_memory=True)
                c_iter = iter(c_dataloader)
                for batch_idx in range(num_samples//batch_size):
                    real_images, real_labels = next(c_iter)
                    real_images = ops.quantize_images(real_images)
                    real_images = ops.resize_images(real_images, resizer, totensor, mean, std, self.local_rank)
                    real_embed = torch.squeeze(resnet50_conv(real_images))
                    if batch_idx == 0:
                        distances = torch.square(real_embed - fake_anchor_embed).mean(dim=1).detach().cpu().numpy()
                        image_holder = real_images.detach().cpu().numpy()
                    else:
                        distances = np.concatenate([
                            distances,
                            torch.square(real_embed - fake_anchor_embed).mean(dim=1).detach().cpu().numpy()
                        ],
                                                   axis=0)
                        image_holder = np.concatenate([image_holder, real_images.detach().cpu().numpy()], axis=0)

                nearest_indices = (-distances).argsort()[-(num_cols - 1):][::-1]
                if c % num_rows == 0:
                    canvas = np.concatenate([fake_anchor.detach().cpu().numpy(), image_holder[nearest_indices]], axis=0)
                elif c % num_rows == num_rows - 1:
                    row_images = np.concatenate([fake_anchor.detach().cpu().numpy(), image_holder[nearest_indices]], axis=0)
                    canvas = np.concatenate((canvas, row_images), axis=0)
                    misc.plot_img_canvas(images=torch.from_numpy(canvas),
                                         save_path=join(self.RUN.save_dir, "figures/{run_name}/fake_anchor_{num_cols}NN_{cls}_classes.png".\
                                                        format(run_name=self.run_name, num_cols=num_cols, cls=c+1)),
                                         num_cols=num_cols,
                                         logger=self.logger,
                                         logging=self.global_rank == 0 and self.logger)
                else:
                    row_images = np.concatenate([fake_anchor.detach().cpu().numpy(), image_holder[nearest_indices]], axis=0)
                    canvas = np.concatenate((canvas, row_images), axis=0)

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # conduct latent interpolation analysis to identify the quaility of latent space (Z)
    # -----------------------------------------------------------------------------
    def run_linear_interpolation(self, num_rows, num_cols, fix_z, fix_y, num_saves=100):
        assert int(fix_z) * int(fix_y) != 1, "unable to switch fix_z and fix_y on together!"
        if self.global_rank == 0:
            flag = "fix_z" if fix_z else "fix_y"
            self.logger.info("Run linear interpolation analysis ({flag}) {num} times.".format(flag=flag, num=num_saves))
        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            shared = misc.peel_model(generator).shared
            for ns in tqdm(range(num_saves)):
                if fix_z:
                    zs = torch.randn(num_rows, 1, self.MODEL.z_dim, device=self.local_rank)
                    zs = zs.repeat(1, num_cols, 1).view(-1, self.MODEL.z_dim)
                    name = "fix_z"
                else:
                    zs = misc.interpolate(torch.randn(num_rows, 1, self.MODEL.z_dim, device=self.local_rank),
                                          torch.randn(num_rows, 1, self.MODEL.z_dim, device=self.local_rank),
                                          num_cols - 2).view(-1, self.MODEL.z_dim)

                if fix_y:
                    ys = sample.sample_onehot(batch_size=num_rows,
                                              num_classes=self.DATA.num_classes,
                                              device=self.local_rank)
                    ys = shared(ys).view(num_rows, 1, -1)
                    ys = ys.repeat(1, num_cols, 1).view(num_rows * (num_cols), -1)
                    name = "fix_y"
                else:
                    ys = misc.interpolate(
                        shared(sample.sample_onehot(num_rows, self.DATA.num_classes)).view(num_rows, 1, -1),
                        shared(sample.sample_onehot(num_rows, self.DATA.num_classes)).view(num_rows, 1, -1),
                        num_cols - 2).view(num_rows * (num_cols), -1)

                interpolated_images = generator(zs, None, shared_label=ys)

                misc.plot_img_canvas(images=interpolated_images.detach().cpu(),
                                     save_path=join(self.RUN.save_dir, "figures/{run_name}/{num}_Interpolated_images_{fix_flag}.png".\
                                                    format(num=ns, run_name=self.run_name, fix_flag=name)),
                                     num_cols=num_cols,
                                     logger=self.logger,
                                     logging=False)

        if self.global_rank == 0 and self.logger:
            print("Save figures to {}/*_Interpolated_images_{}.png".format(
                join(self.RUN.save_dir, "figures", self.run_name), flag))

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # visualize shifted fourier spectrums of real and fake images
    # -----------------------------------------------------------------------------
    def run_frequency_analysis(self, dataloader):
        if self.global_rank == 0:
            self.logger.info("Run frequency analysis (use {num} fake and {ref} images).".\
                             format(num=len(dataloader), ref=self.RUN.ref_dataset))
        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            data_iter = iter(dataloader)
            num_batches = len(dataloader) // self.OPTIMIZATION.batch_size
            for i in range(num_batches):
                real_images, real_labels = next(data_iter)
                fake_images, fake_labels, _, _, _, _, _ = sample.generate_images(z_prior=self.MODEL.z_prior,
                                                                          truncation_factor=self.RUN.truncation_factor,
                                                                          batch_size=self.OPTIMIZATION.batch_size,
                                                                          z_dim=self.MODEL.z_dim,
                                                                          num_classes=self.DATA.num_classes,
                                                                          y_sampler="totally_random",
                                                                          radius="N/A",
                                                                          generator=generator,
                                                                          discriminator=self.Dis,
                                                                          is_train=False,
                                                                          LOSS=self.LOSS,
                                                                          RUN=self.RUN,
                                                                          MODEL=self.MODEL,
                                                                          device=self.local_rank,
                                                                          is_stylegan=self.is_stylegan,
                                                                          generator_mapping=generator_mapping,
                                                                          generator_synthesis=generator_synthesis,
                                                                          style_mixing_p=0.0,
                                                                          stylegan_update_emas=False,
                                                                          cal_trsp_cost=False)
                fake_images = fake_images.detach().cpu().numpy()

                real_images = np.asarray((real_images + 1) * 127.5, np.uint8)
                fake_images = np.asarray((fake_images + 1) * 127.5, np.uint8)

                if i == 0:
                    real_array = real_images
                    fake_array = fake_images
                else:
                    real_array = np.concatenate([real_array, real_images], axis=0)
                    fake_array = np.concatenate([fake_array, fake_images], axis=0)

            N, C, H, W = np.shape(real_array)
            real_r, real_g, real_b = real_array[:, 0, :, :], real_array[:, 1, :, :], real_array[:, 2, :, :]
            real_gray = 0.2989 * real_r + 0.5870 * real_g + 0.1140 * real_b
            fake_r, fake_g, fake_b = fake_array[:, 0, :, :], fake_array[:, 1, :, :], fake_array[:, 2, :, :]
            fake_gray = 0.2989 * fake_r + 0.5870 * fake_g + 0.1140 * fake_b
            for j in tqdm(range(N)):
                real_gray_f = np.fft.fft2(real_gray[j] - ndimage.median_filter(real_gray[j], size=H // 8))
                fake_gray_f = np.fft.fft2(fake_gray[j] - ndimage.median_filter(fake_gray[j], size=H // 8))

                real_gray_f_shifted = np.fft.fftshift(real_gray_f)
                fake_gray_f_shifted = np.fft.fftshift(fake_gray_f)

                if j == 0:
                    real_gray_spectrum = 20 * np.log(np.abs(real_gray_f_shifted)) / N
                    fake_gray_spectrum = 20 * np.log(np.abs(fake_gray_f_shifted)) / N
                else:
                    real_gray_spectrum += 20 * np.log(np.abs(real_gray_f_shifted)) / N
                    fake_gray_spectrum += 20 * np.log(np.abs(fake_gray_f_shifted)) / N

        misc.plot_spectrum_image(real_spectrum=real_gray_spectrum,
                                 fake_spectrum=fake_gray_spectrum,
                                 directory=join(self.RUN.save_dir, "figures", self.run_name),
                                 logger=self.logger,
                                 logging=self.global_rank == 0 and self.logger)

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # visualize discriminator's embeddings of real or fake images using TSNE
    # -----------------------------------------------------------------------------
    def run_tsne(self, dataloader):
        if self.global_rank == 0:
            self.logger.info("Start TSNE analysis using randomly sampled 10 classes.")
            self.logger.info("Use {ref} dataset and the same amount of generated images for visualization.".format(
                ref=self.RUN.ref_dataset))
        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            save_output, real, fake, hook_handles = misc.SaveOutput(), {}, {}, []
            for name, layer in misc.peel_model(self.Dis).named_children():
                if name == "linear1":
                    handle = layer.register_forward_pre_hook(save_output)
                    hook_handles.append(handle)

            tsne_iter = iter(dataloader)
            num_batches = len(dataloader.dataset) // self.OPTIMIZATION.batch_size
            for i in range(num_batches):
                real_images, real_labels = next(tsne_iter)
                real_images, real_labels = real_images.to(self.local_rank), real_labels.to(self.local_rank)

                real_dict = self.Dis(real_images, real_labels)
                if i == 0:
                    real["embeds"] = save_output.outputs[0][0].detach().cpu().numpy()
                    real["labels"] = real_labels.detach().cpu().numpy()
                else:
                    real["embeds"] = np.concatenate([real["embeds"], save_output.outputs[0][0].cpu().detach().numpy()],
                                                    axis=0)
                    real["labels"] = np.concatenate([real["labels"], real_labels.detach().cpu().numpy()])

                save_output.clear()

                fake_images, fake_labels, _, _, _, _, _ = sample.generate_images(z_prior=self.MODEL.z_prior,
                                                                           truncation_factor=self.RUN.truncation_factor,
                                                                           batch_size=self.OPTIMIZATION.batch_size,
                                                                           z_dim=self.MODEL.z_dim,
                                                                           num_classes=self.DATA.num_classes,
                                                                           y_sampler="totally_random",
                                                                           radius="N/A",
                                                                           generator=generator,
                                                                           discriminator=self.Dis,
                                                                           is_train=False,
                                                                           LOSS=self.LOSS,
                                                                           RUN=self.RUN,
                                                                           MODEL=self.MODEL,
                                                                           device=self.local_rank,
                                                                           is_stylegan=self.is_stylegan,
                                                                           generator_mapping=generator_mapping,
                                                                           generator_synthesis=generator_synthesis,
                                                                           style_mixing_p=0.0,
                                                                           stylegan_update_emas=False,
                                                                           cal_trsp_cost=False)

                fake_dict = self.Dis(fake_images, fake_labels)
                if i == 0:
                    fake["embeds"] = save_output.outputs[0][0].detach().cpu().numpy()
                    fake["labels"] = fake_labels.detach().cpu().numpy()
                else:
                    fake["embeds"] = np.concatenate([fake["embeds"], save_output.outputs[0][0].cpu().detach().numpy()],
                                                    axis=0)
                    fake["labels"] = np.concatenate([fake["labels"], fake_labels.detach().cpu().numpy()])

                save_output.clear()

            tsne = TSNE(n_components=2, verbose=1, perplexity=40, n_iter=300)
            if self.DATA.num_classes > 10:
                cls_indices = np.random.permutation(self.DATA.num_classes)[:10]
                real["embeds"] = real["embeds"][np.isin(real["labels"], cls_indices)]
                real["labels"] = real["labels"][np.isin(real["labels"], cls_indices)]
                fake["embeds"] = fake["embeds"][np.isin(fake["labels"], cls_indices)]
                fake["labels"] = fake["labels"][np.isin(fake["labels"], cls_indices)]

            real_tsne_results = tsne.fit_transform(real["embeds"])
            misc.plot_tsne_scatter_plot(df=real,
                                        tsne_results=real_tsne_results,
                                        flag="real",
                                        directory=join(self.RUN.save_dir, "figures", self.run_name),
                                        logger=self.logger,
                                        logging=self.global_rank == 0 and self.logger)

            fake_tsne_results = tsne.fit_transform(fake["embeds"])
            misc.plot_tsne_scatter_plot(df=fake,
                                        tsne_results=fake_tsne_results,
                                        flag="fake",
                                        directory=join(self.RUN.save_dir, "figures", self.run_name),
                                        logger=self.logger,
                                        logging=self.global_rank == 0 and self.logger)

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # calculate intra-class FID (iFID) to identify intra-class diversity
    # -----------------------------------------------------------------------------
    def calculate_intra_class_fid(self, dataset):
        if self.global_rank == 0:
            self.logger.info("Start calculating iFID (use approx. {num} fake images per class and train images as the reference).".\
                             format(num=int(len(dataset)/self.DATA.num_classes)))

        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        fids = []
        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            for c in tqdm(range(self.DATA.num_classes)):
                num_samples, target_sampler = sample.make_target_cls_sampler(dataset, c)
                batch_size = self.OPTIMIZATION.batch_size if num_samples >= self.OPTIMIZATION.batch_size else num_samples
                dataloader = torch.utils.data.DataLoader(dataset,
                                                         batch_size=batch_size,
                                                         shuffle=False,
                                                         sampler=target_sampler,
                                                         num_workers=self.RUN.num_workers,
                                                         pin_memory=True,
                                                         drop_last=False)

                mu, sigma = fid.calculate_moments(data_loader=dataloader,
                                                  eval_model=self.eval_model,
                                                  num_generate="N/A",
                                                  batch_size=batch_size,
                                                  quantize=True,
                                                  world_size=self.OPTIMIZATION.world_size,
                                                  DDP=self.DDP,
                                                  disable_tqdm=True,
                                                  fake_feats=None)

                c_fake_feats, _,_ = features.generate_images_and_stack_features(
                                                    generator=generator,
                                                    discriminator=self.Dis,
                                                    eval_model=self.eval_model,
                                                    num_generate=num_samples,
                                                    y_sampler=c,
                                                    batch_size=self.OPTIMIZATION.batch_size,
                                                    z_prior=self.MODEL.z_prior,
                                                    truncation_factor=self.RUN.truncation_factor,
                                                    z_dim=self.MODEL.z_dim,
                                                    num_classes=self.DATA.num_classes,
                                                    LOSS=self.LOSS,
                                                    RUN=self.RUN,
                                                    MODEL=self.MODEL,
                                                    is_stylegan=self.is_stylegan,
                                                    generator_mapping=generator_mapping,
                                                    generator_synthesis=generator_synthesis,
                                                    quantize=True,
                                                    world_size=self.OPTIMIZATION.world_size,
                                                    DDP=self.DDP,
                                                    device=self.local_rank,
                                                    logger=self.logger,
                                                    disable_tqdm=True)

                ifid_score, _, _ = fid.calculate_fid(data_loader="N/A",
                                                     eval_model=self.eval_model,
                                                     num_generate=num_samples,
                                                     cfgs=self.cfgs,
                                                     pre_cal_mean=mu,
                                                     pre_cal_std=sigma,
                                                     quantize=False,
                                                     fake_feats=c_fake_feats,
                                                     disable_tqdm=True)

                fids.append(ifid_score)

                # save iFID values in .npz format
                metric_dict = {"iFID": ifid_score}

                save_dict = misc.accm_values_convert_dict(list_dict={"iFID": []},
                                                          value_dict=metric_dict,
                                                          step=c,
                                                          interval=1)
                misc.save_dict_npy(directory=join(self.RUN.save_dir, "statistics", self.run_name),
                                   name="iFID",
                                   dictionary=save_dict)

        if self.global_rank == 0 and self.logger:
            self.logger.info("Average iFID score: {iFID}".format(iFID=sum(fids, 0.0) / len(fids)))

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # perform semantic (closed-form) factorization for latent nevigation
    # -----------------------------------------------------------------------------
    def run_semantic_factorization(self, num_rows, num_cols, maximum_variations):
        if self.global_rank == 0:
            self.logger.info("Perform semantic factorization for latent nevigation.")

        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        requires_grad = self.LOSS.apply_lo or self.RUN.langevin_sampling
        with torch.no_grad() if not requires_grad else misc.dummy_context_mgr() as ctx:
            misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
            generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

            zs, fake_labels, _ = sample.sample_zy(z_prior=self.MODEL.z_prior,
                                                  batch_size=self.OPTIMIZATION.batch_size,
                                                  z_dim=self.MODEL.z_dim,
                                                  num_classes=self.DATA.num_classes,
                                                  truncation_factor=self.RUN.truncation_factor,
                                                  y_sampler="totally_random",
                                                  radius="N/A",
                                                  device=self.local_rank)

            for i in tqdm(range(self.OPTIMIZATION.batch_size)):
                images_canvas = sefa.apply_sefa(generator=generator,
                                                backbone=self.MODEL.backbone,
                                                z=zs[i],
                                                fake_label=fake_labels[i],
                                                num_semantic_axis=num_rows,
                                                maximum_variations=maximum_variations,
                                                num_cols=num_cols)

                misc.plot_img_canvas(images=images_canvas.detach().cpu(),
                                     save_path=join(self.RUN.save_dir, "figures/{run_name}/{idx}_sefa_images.png".\
                                                    format(idx=i, run_name=self.run_name)),
                                     num_cols=num_cols,
                                     logger=self.logger,
                                     logging=False)

        if self.global_rank == 0 and self.logger:
            print("Save figures to {}/*_sefa_images.png".format(join(self.RUN.save_dir, "figures", self.run_name)))

        misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis)

    # -----------------------------------------------------------------------------
    # compute classifier accuracy score (CAS) to identify class-conditional precision and recall
    # -----------------------------------------------------------------------------
    def compute_GAN_train_or_test_classifier_accuracy_score(self, GAN_train=False, GAN_test=False):
        assert GAN_train*GAN_test == 0, "cannot conduct GAN_train and GAN_test togather."
        if self.global_rank == 0:
            if GAN_train:
                phase, metric = "train", "recall"
            else:
                phase, metric = "test", "precision"
            self.logger.info("compute GAN_{phase} Classifier Accuracy Score (CAS) to identify class-conditional {metric}.". \
                             format(phase=phase, metric=metric))

        if self.gen_ctlr.standing_statistics:
            self.gen_ctlr.std_stat_counter += 1

        misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis)
        generator, generator_mapping, generator_synthesis = self.gen_ctlr.prepare_generator()

        best_top1, best_top5, cas_setting = 0.0, 0.0, self.MISC.cas_setting[self.DATA.name]
        model = resnet.ResNet(dataset=self.DATA.name,
                              depth=cas_setting["depth"],
                              num_classes=self.DATA.num_classes,
                              bottleneck=cas_setting["bottleneck"]).to("cuda")

        optimizer = torch.optim.SGD(params=model.parameters(),
                                    lr=cas_setting["lr"],
                                    momentum=cas_setting["momentum"],
                                    weight_decay=cas_setting["weight_decay"],
                                    nesterov=True)

        if self.OPTIMIZATION.world_size > 1:
            model = DataParallel(model, output_device=self.local_rank)

        epoch_trained = 0
        if self.RUN.ckpt_dir is not None and self.RUN.resume_classifier_train:
            is_pre_trained_model, mode = ckpt.check_is_pre_trained_model(ckpt_dir=self.RUN.ckpt_dir,
                                                                         GAN_train=GAN_train,
                                                                         GAN_test=GAN_test)
            if is_pre_trained_model:
                epoch_trained, best_top1, best_top5, best_epoch = ckpt.load_GAN_train_test_model(model=model,
                                                                                                 mode=mode,
                                                                                                 optimizer=optimizer,
                                                                                                 RUN=self.RUN)

        for current_epoch in tqdm(range(epoch_trained, cas_setting["epochs"])):
            model.train()
            optimizer.zero_grad()
            ops.adjust_learning_rate(optimizer=optimizer,
                                     lr_org=cas_setting["lr"],
                                     epoch=current_epoch,
                                     total_epoch=cas_setting["epochs"],
                                     dataset=self.DATA.name)

            train_top1_acc, train_top5_acc, train_loss = misc.AverageMeter(), misc.AverageMeter(), misc.AverageMeter()
            for i, (images, labels) in enumerate(self.train_dataloader):
                if GAN_train:
                    images, labels, _, _, _, _, _ = sample.generate_images(z_prior=self.MODEL.z_prior,
                                                                     truncation_factor=self.RUN.truncation_factor,
                                                                     batch_size=self.OPTIMIZATION.batch_size,
                                                                     z_dim=self.MODEL.z_dim,
                                                                     num_classes=self.DATA.num_classes,
                                                                     y_sampler="totally_random",
                                                                     radius="N/A",
                                                                     generator=generator,
                                                                     discriminator=self.Dis,
                                                                     is_train=False,
                                                                     LOSS=self.LOSS,
                                                                     RUN=self.RUN,
                                                                     MODEL=self.MODEL,
                                                                     device=self.local_rank,
                                                                     is_stylegan=self.is_stylegan,
                                                                     generator_mapping=generator_mapping,
                                                                     generator_synthesis=generator_synthesis,
                                                                     style_mixing_p=0.0,
                                                                     stylegan_update_emas=False,
                                                                     cal_trsp_cost=False)
                else:
                    images, labels = images.to(self.local_rank), labels.to(self.local_rank)

                logits = model(images)
                ce_loss = self.ce_loss(logits, labels)

                train_acc1, train_acc5 = misc.accuracy(logits.data, labels, topk=(1, 5))

                train_loss.update(ce_loss.item(), images.size(0))
                train_top1_acc.update(train_acc1.item(), images.size(0))
                train_top5_acc.update(train_acc5.item(), images.size(0))

                ce_loss.backward()
                optimizer.step()

            valid_acc1, valid_acc5, valid_loss = self.validate_classifier(model=model,
                                                                          generator=generator,
                                                                          generator_mapping=generator_mapping,
                                                                          generator_synthesis=generator_synthesis,
                                                                          epoch=current_epoch,
                                                                          GAN_test=GAN_test,
                                                                          setting=cas_setting)

            is_best = valid_acc1 > best_top1
            best_top1 = max(valid_acc1, best_top1)
            if is_best:
                best_top5, best_epoch = valid_acc5, current_epoch
                model_ = misc.peel_model(model)
                states = {"state_dict": model_.state_dict(), "optimizer": optimizer.state_dict(), "epoch": current_epoch+1,
                          "best_top1": best_top1, "best_top5": best_top5, "best_epoch": best_epoch}
                misc.save_model_c(states, mode, self.RUN)

            if self.local_rank == 0:
                self.logger.info("Current best accuracy: Top-1: {top1:.4f}% and Top-5 {top5:.4f}%".format(top1=best_top1, top5=best_top5))
                self.logger.info("Save model to {}".format(self.RUN.ckpt_dir))

    # -----------------------------------------------------------------------------
    # validate GAN_train or GAN_test classifier using generated or training dataset
    # -----------------------------------------------------------------------------
    def validate_classifier(self,model, generator, generator_mapping, generator_synthesis, epoch, GAN_test, setting):
        model.eval()
        valid_top1_acc, valid_top5_acc, valid_loss = misc.AverageMeter(), misc.AverageMeter(), misc.AverageMeter()
        for i, (images, labels) in enumerate(self.train_dataloader):
            if GAN_test:
                images, labels, _, _, _, _, _ = sample.generate_images(z_prior=self.MODEL.z_prior,
                                                                 truncation_factor=self.RUN.truncation_factor,
                                                                 batch_size=self.OPTIMIZATION.batch_size,
                                                                 z_dim=self.MODEL.z_dim,
                                                                 num_classes=self.DATA.num_classes,
                                                                 y_sampler="totally_random",
                                                                 radius="N/A",
                                                                 generator=generator,
                                                                 discriminator=self.Dis,
                                                                 is_train=False,
                                                                 LOSS=self.LOSS,
                                                                 RUN=self.RUN,
                                                                 MODEL=self.MODEL,
                                                                 device=self.local_rank,
                                                                 is_stylegan=self.is_stylegan,
                                                                 generator_mapping=generator_mapping,
                                                                 generator_synthesis=generator_synthesis,
                                                                 style_mixing_p=0.0,
                                                                 stylegan_update_emas=False,
                                                                 cal_trsp_cost=False)
            else:
                images, labels = images.to(self.local_rank), labels.to(self.local_rank)

            output = model(images)
            ce_loss = self.ce_loss(output, labels)

            valid_acc1, valid_acc5 = misc.accuracy(output.data, labels, topk=(1, 5))

            valid_loss.update(ce_loss.item(), images.size(0))
            valid_top1_acc.update(valid_acc1.item(), images.size(0))
            valid_top5_acc.update(valid_acc5.item(), images.size(0))

        if self.local_rank == 0:
            self.logger.info("Top 1-acc {top1.val:.4f} ({top1.avg:.4f})\t"
                             "Top 5-acc {top5.val:.4f} ({top5.avg:.4f})".format(top1=valid_top1_acc, top5=valid_top5_acc))
        return valid_top1_acc.avg, valid_top5_acc.avg, valid_loss.avg