code/studiogan/src/models/model.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/models/model.py

import copy

from torch.nn import DataParallel
from torch.nn.parallel import DistributedDataParallel as DDP
import torch

from sync_batchnorm.batchnorm import convert_model
from utils.ema import Ema
from utils.ema import EmaStylegan2
import utils.misc as misc


def load_generator_discriminator(DATA, OPTIMIZATION, MODEL, STYLEGAN, MODULES, RUN, device, logger):
    if device == 0:
        logger.info("Build a Generative Adversarial Network.")
    module = __import__("models.{backbone}".format(backbone=MODEL.backbone), fromlist=["something"])
    if device == 0:
        logger.info("Modules are located on './src/models.{backbone}'.".format(backbone=MODEL.backbone))

    if MODEL.backbone in ["stylegan2", "stylegan3"]:
        channel_base, channel_max = 32768 if MODEL.backbone == "stylegan3" or DATA.img_size >= 512 or \
                                    DATA.name in ["CIFAR10", "CIFAR100"] else 16384, 512
        gen_c_dim = DATA.num_classes if MODEL.g_cond_mtd == "cAdaIN" else 0
        dis_c_dim = DATA.num_classes if MODEL.d_cond_mtd in STYLEGAN.cond_type else 0
        if RUN.mixed_precision:
            num_fp16_res = 4
            conv_clamp = 256
        else:
            num_fp16_res = 0
            conv_clamp = None
        if MODEL.backbone == "stylegan2":
            mapping_kwargs = {"num_layers": STYLEGAN.mapping_network}
            if getattr(STYLEGAN, "use_rtm_mapper", False):
                mapping_kwargs.update({
                    "use_rtm_mapper": True,
                    "rtm_num_tokens": STYLEGAN.rtm_num_tokens,
                    "rtm_H_cycles": STYLEGAN.rtm_H_cycles,
                    "rtm_L_cycles": STYLEGAN.rtm_L_cycles,
                    "rtm_H_layers": STYLEGAN.rtm_H_layers,
                    "rtm_L_layers": STYLEGAN.rtm_L_layers,
                    "rtm_hidden_size": STYLEGAN.rtm_hidden_size,
                    "rtm_expansion": STYLEGAN.rtm_expansion,
                    "rtm_refinement_steps": STYLEGAN.rtm_refinement_steps,
                    "rtm_with_grad": STYLEGAN.rtm_with_grad,
                    "rtm_cycle_noise_std": STYLEGAN.rtm_cycle_noise_std,
                    "use_rtm_equalized": getattr(STYLEGAN, "use_rtm_equalized", False),
                    "rtm_lr_multiplier": getattr(STYLEGAN, "rtm_lr_multiplier", 1.0),
                })
            Gen = module.Generator(z_dim=MODEL.z_dim,
                                c_dim=gen_c_dim,
                                w_dim=MODEL.w_dim,
                                img_resolution=DATA.img_size,
                                img_channels=DATA.img_channels,
                                MODEL=MODEL,
                                mapping_kwargs=mapping_kwargs,
                                synthesis_kwargs={"channel_base": channel_base, "channel_max": channel_max, \
                                "num_fp16_res": num_fp16_res, "conv_clamp": conv_clamp}).to(device)
        else:
            magnitude_ema_beta = 0.5 ** (OPTIMIZATION.batch_size * OPTIMIZATION.acml_steps / (20 * 1e3))
            g_channel_base, g_channel_max, conv_kernel, use_radial_filters = channel_base, channel_max, 3, False
            if STYLEGAN.stylegan3_cfg == "stylegan3-r":
                g_channel_base, g_channel_max, conv_kernel, use_radial_filters = channel_base * 2, channel_max * 2, 1, True
            Gen = module.Generator(z_dim=MODEL.z_dim,
                                c_dim=gen_c_dim,
                                w_dim=MODEL.w_dim,
                                img_resolution=DATA.img_size,
                                img_channels=DATA.img_channels,
                                MODEL=MODEL,
                                mapping_kwargs={"num_layers": STYLEGAN.mapping_network},
                                synthesis_kwargs={"channel_base": g_channel_base, "channel_max": g_channel_max, \
                                "num_fp16_res": num_fp16_res, "conv_clamp": conv_clamp, "conv_kernel": conv_kernel, \
                                "use_radial_filters": use_radial_filters, "magnitude_ema_beta": magnitude_ema_beta}).to(device)

        Gen_mapping, Gen_synthesis = Gen.mapping, Gen.synthesis

        module = __import__("models.stylegan2", fromlist=["something"]) # always use StyleGAN2 discriminator
        Dis = module.Discriminator(c_dim=dis_c_dim,
                                   img_resolution=DATA.img_size,
                                   img_channels=DATA.img_channels,
                                   architecture=STYLEGAN.d_architecture,
                                   channel_base=channel_base,
                                   channel_max=channel_max,
                                   num_fp16_res=num_fp16_res,
                                   conv_clamp=conv_clamp,
                                   cmap_dim=None,
                                   d_cond_mtd=MODEL.d_cond_mtd,
                                   aux_cls_type=MODEL.aux_cls_type,
                                   d_embed_dim=MODEL.d_embed_dim,
                                   num_classes=DATA.num_classes,
                                   normalize_d_embed=MODEL.normalize_d_embed,
                                   block_kwargs={},
                                   mapping_kwargs={},
                                   epilogue_kwargs={
                                       "mbstd_group_size": STYLEGAN.d_epilogue_mbstd_group_size
                                   },
                                   MODEL=MODEL).to(device)

        if MODEL.apply_g_ema:
            if device == 0:
                logger.info("Prepare exponential moving average generator with decay rate of {decay}."\
                            .format(decay=MODEL.g_ema_decay))
            Gen_ema = copy.deepcopy(Gen)
            Gen_ema_mapping, Gen_ema_synthesis = Gen_ema.mapping, Gen_ema.synthesis

            ema = EmaStylegan2(source=Gen,
                               target=Gen_ema,
                               ema_kimg=STYLEGAN.g_ema_kimg,
                               ema_rampup=STYLEGAN.g_ema_rampup,
                               effective_batch_size=OPTIMIZATION.batch_size * OPTIMIZATION.acml_steps)
        else:
            Gen_ema, Gen_ema_mapping, Gen_ema_synthesis, ema = None, None, None, None

    else:
        Gen = module.Generator(z_dim=MODEL.z_dim,
                               g_shared_dim=MODEL.g_shared_dim,
                               img_size=DATA.img_size,
                               g_conv_dim=MODEL.g_conv_dim,
                               apply_attn=MODEL.apply_attn,
                               attn_g_loc=MODEL.attn_g_loc,
                               g_cond_mtd=MODEL.g_cond_mtd,
                               num_classes=DATA.num_classes,
                               g_init=MODEL.g_init,
                               g_depth=MODEL.g_depth,
                               mixed_precision=RUN.mixed_precision,
                               MODULES=MODULES,
                               MODEL=MODEL).to(device)

        Gen_mapping, Gen_synthesis = None, None

        Dis = module.Discriminator(img_size=DATA.img_size,
                                   d_conv_dim=MODEL.d_conv_dim,
                                   apply_d_sn=MODEL.apply_d_sn,
                                   apply_attn=MODEL.apply_attn,
                                   attn_d_loc=MODEL.attn_d_loc,
                                   d_cond_mtd=MODEL.d_cond_mtd,
                                   aux_cls_type=MODEL.aux_cls_type,
                                   d_embed_dim=MODEL.d_embed_dim,
                                   num_classes=DATA.num_classes,
                                   normalize_d_embed=MODEL.normalize_d_embed,
                                   d_init=MODEL.d_init,
                                   d_depth=MODEL.d_depth,
                                   mixed_precision=RUN.mixed_precision,
                                   MODULES=MODULES,
                                   MODEL=MODEL).to(device)
        if MODEL.apply_g_ema:
            if device == 0:
                logger.info("Prepare exponential moving average generator with decay rate of {decay}."\
                            .format(decay=MODEL.g_ema_decay))
            Gen_ema = copy.deepcopy(Gen)
            Gen_ema_mapping, Gen_ema_synthesis = None, None

            ema = Ema(source=Gen, target=Gen_ema, decay=MODEL.g_ema_decay, start_iter=MODEL.g_ema_start)
        else:
            Gen_ema, Gen_ema_mapping, Gen_ema_synthesis, ema = None, None, None, None

    if device == 0:
        logger.info(misc.count_parameters(Gen))
    if device == 0:
        logger.info(Gen)

    if device == 0:
        logger.info(misc.count_parameters(Dis))
    if device == 0:
        logger.info(Dis)
    return Gen, Gen_mapping, Gen_synthesis, Dis, Gen_ema, Gen_ema_mapping, Gen_ema_synthesis, ema


def prepare_parallel_training(Gen, Gen_mapping, Gen_synthesis, Dis, Gen_ema, Gen_ema_mapping, Gen_ema_synthesis,
                              MODEL, world_size, distributed_data_parallel, synchronized_bn, apply_g_ema, device,
                              STYLEGAN=None):
    if distributed_data_parallel:
        if synchronized_bn:
            process_group = torch.distributed.new_group([w for w in range(world_size)])
            Gen = torch.nn.SyncBatchNorm.convert_sync_batchnorm(Gen, process_group)
            Dis = torch.nn.SyncBatchNorm.convert_sync_batchnorm(Dis, process_group)
            if apply_g_ema:
                Gen_ema = torch.nn.SyncBatchNorm.convert_sync_batchnorm(Gen_ema, process_group)

        # The RTM mapper runs (H_cycles - 1) of its inner H/L iterations under
        # torch.no_grad() when rtm_with_grad is False, so the H_init / L_init
        # initial-state Parameters do not always receive gradients. DDP rejects
        # that unless we explicitly allow unused parameters on the mapping.
        gen_mapping_find_unused = bool(getattr(STYLEGAN, "use_rtm_mapper", False))

        if MODEL.backbone in ["stylegan2", "stylegan3"]:
            Gen_mapping = DDP(Gen.mapping, device_ids=[device], broadcast_buffers=False,
                              find_unused_parameters=gen_mapping_find_unused)
            Gen_synthesis = DDP(Gen.synthesis, device_ids=[device], broadcast_buffers=False)
        else:
            Gen = DDP(Gen, device_ids=[device], broadcast_buffers=synchronized_bn)
        Dis = DDP(Dis, device_ids=[device],
                  broadcast_buffers=False if MODEL.backbone in ["stylegan2", "stylegan3"] else synchronized_bn,
                  find_unused_parameters=True if MODEL.info_type in ["discrete", "continuous", "both"] else False)
        if apply_g_ema:
            if MODEL.backbone in ["stylegan2", "stylegan3"]:
                Gen_ema_mapping = DDP(Gen_ema.mapping, device_ids=[device], broadcast_buffers=False,
                                      find_unused_parameters=gen_mapping_find_unused)
                Gen_ema_synthesis = DDP(Gen_ema.synthesis, device_ids=[device], broadcast_buffers=False)
            else:
                Gen_ema = DDP(Gen_ema, device_ids=[device], broadcast_buffers=synchronized_bn)
    else:
        if MODEL.backbone in ["stylegan2", "stylegan3"]:
            Gen_mapping = DataParallel(Gen.mapping, output_device=device)
            Gen_synthesis = DataParallel(Gen.synthesis, output_device=device)
        else:
            Gen = DataParallel(Gen, output_device=device)
        Dis = DataParallel(Dis, output_device=device)
        if apply_g_ema:
            if MODEL.backbone in ["stylegan2", "stylegan3"]:
                Gen_ema_mapping = DataParallel(Gen_ema.mapping, output_device=device)
                Gen_ema_synthesis = DataParallel(Gen_ema.synthesis, output_device=device)
            else:
                Gen_ema = DataParallel(Gen_ema, output_device=device)

        if synchronized_bn:
            Gen = convert_model(Gen).to(device)
            Dis = convert_model(Dis).to(device)
            if apply_g_ema:
                Gen_ema = convert_model(Gen_ema).to(device)
    return Gen, Gen_mapping, Gen_synthesis, Dis, Gen_ema, Gen_ema_mapping, Gen_ema_synthesis