train_profile.py

import argparse
import gc
import math
import os
import random
import sys
import time

import numpy as np
import torch as th
import torch.autograd.profiler as profiler
import wandb
from contrastive_learner import ContrastiveLearner, RandomApply
from kornia import augmentation as augs
from scipy.ndimage import gaussian_filter
from torch import nn
from torch.nn import functional as F
from torch.utils import data
from torchvision import transforms, utils
from tqdm import tqdm

import validation
from augment import augment
from dataset import MultiResolutionDataset
from distributed import get_rank, reduce_loss_dict, reduce_sum, synchronize
from lookahead_minimax import LookaheadMinimax
from models.stylegan2 import Discriminator, Generator

sys.path.insert(0, "../lookahead_minimax")


def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle:
        return data.RandomSampler(dataset)
    else:
        return data.SequentialSampler(dataset)


def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag


def accumulate(model1, model2, decay=0.5 ** (32.0 / 10_000)):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())
    for name, param in model1.named_parameters():
        param.data = decay * par1[name].data + (1 - decay) * par2[name].data


def sample_data(loader):
    while True:
        for batch in loader:
            yield batch


def make_noise(batch_size, latent_dim, prob):
    if prob > 0 and random.random() < prob:
        return th.randn(2, batch_size, latent_dim, device=device).unbind(0)
    else:
        return [th.randn(batch_size, latent_dim, device=device)]


def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)
    return real_loss.mean() + fake_loss.mean()


def d_r1_penalty(real_img, real_pred, args):
    (grad_real,) = th.autograd.grad(real_pred.sum(), real_img, create_graph=True)
    r1_loss = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()
    r1_loss = r1_loss / 2.0 + 0 * real_pred[0]
    return r1_loss


def g_non_saturating_loss(fake_pred):
    return F.softplus(-fake_pred).mean()


def g_path_length_regularization(generator, mean_path_length, args):
    path_batch_size = max(1, args.batch_size // args.path_batch_shrink)

    noise = make_noise(path_batch_size, args.latent_size, args.mixing_prob)
    fake_img, latents = generator(noise, return_latents=True)

    img_noise = th.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])
    noisy_img_sum = (fake_img * img_noise).sum()

    (grad,) = th.autograd.grad(noisy_img_sum, latents, create_graph=True)

    path_lengths = th.sqrt(grad.pow(2).sum(2).mean(1))
    path_mean = mean_path_length + 0.01 * (path_lengths.mean() - mean_path_length)
    path_loss = (path_lengths - path_mean).pow(2).mean()
    if not th.isnan(path_mean):
        mean_path_length = path_mean.detach()

    if args.path_batch_shrink:
        path_loss += 0 * fake_img[0, 0, 0, 0]

    return path_loss, mean_path_length


# detach / item all the things
# separate out into smaller functions so those local scopes get cleaned better
# placeholder tensor pattern
# torch no grad where possible
def train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema):
    if args.distributed:
        g_module = generator.module
        d_module = discriminator.module
        if contrast_learner is not None:
            cl_module = contrast_learner.module
    else:
        g_module = generator
        d_module = discriminator
        cl_module = contrast_learner

    loader = sample_data(loader)
    sample_z = th.randn(args.n_sample, args.latent_size, device=device)
    mse = th.nn.MSELoss()
    mean_path_length = 0
    ada_augment = th.tensor([0.0, 0.0], device=device)
    ada_aug_p = args.augment_p if args.augment_p > 0 else 0.0
    ada_aug_step = args.ada_target / args.ada_length
    r_t_stat = 0
    fids = []

    pbar = range(args.iter)
    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0)

    for idx in pbar:
        i = idx + args.start_iter
        if i > args.iter:
            print("Done!")
            break

        if idx == args.nsys_iter:
            print("Profiling begun at iteration {}".format(idx))
            th.cuda.cudart().cudaProfilerStart()
        if idx >= args.nsys_iter:
            th.cuda.nvtx.range_push("Iter {}".format(idx))

        tick_start = time.time()

        loss_dict = {
            "Generator": th.tensor(0, device=device).float(),
            "Discriminator": th.tensor(0, device=device).float(),
            "Real Score": th.tensor(0, device=device).float(),
            "Fake Score": th.tensor(0, device=device).float(),
            "Contrastive": th.tensor(0, device=device).float(),
            "Consistency": th.tensor(0, device=device).float(),
            "R1 Penalty": th.tensor(0, device=device).float(),
            "Path Length Regularization": th.tensor(0, device=device).float(),
            "Augment": th.tensor(0, device=device).float(),
            "Rt": th.tensor(0, device=device).float(),
        }

        with profiler.record_function("D train"):
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("D train")
            requires_grad(generator, False)
            requires_grad(discriminator, True)

            discriminator.zero_grad()
            for _ in range(args.num_accumulate):
                real_img_og = next(loader).to(device, non_blocking=True)
                noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)
                fake_img_og, _ = generator(noise)
                if args.augment:
                    fake_img, _ = augment(fake_img_og, ada_aug_p)
                    real_img, _ = augment(real_img_og, ada_aug_p)
                else:
                    fake_img = fake_img_og
                    real_img = real_img_og

                fake_pred = discriminator(fake_img)
                real_pred = discriminator(real_img)
                logistic_loss = d_logistic_loss(real_pred, fake_pred)
                loss_dict["Discriminator"] += logistic_loss.detach()
                loss_dict["Real Score"] += real_pred.mean().detach()
                loss_dict["Fake Score"] += fake_pred.mean().detach()
                d_loss = logistic_loss

                if args.contrastive > 0:
                    contrast_learner(fake_img_og, fake_img, accumulate=True)
                    contrast_learner(real_img_og, real_img, accumulate=True)
                    contrast_loss = cl_module.calculate_loss()
                    loss_dict["Contrastive"] += contrast_loss.detach()
                    d_loss += args.contrastive * contrast_loss

                if args.balanced_consistency > 0:
                    consistency_loss = mse(real_pred, discriminator(real_img_og)) + mse(
                        fake_pred, discriminator(fake_img_og)
                    )
                    loss_dict["Consistency"] += consistency_loss.detach()
                    d_loss += args.balanced_consistency * consistency_loss

                d_loss /= args.num_accumulate
                d_loss.backward()
            d_optim.step()
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()

        if args.r1 > 0 and i % args.d_reg_every == 0:
            with profiler.record_function("D reg"):
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_push("D reg")
                discriminator.zero_grad()
                for _ in range(args.num_accumulate):
                    real_img = next(loader).to(device, non_blocking=True)
                    real_img.requires_grad = True
                    real_pred = discriminator(real_img)
                    r1_loss = d_r1_penalty(real_img, real_pred, args)
                    loss_dict["R1 Penalty"] += r1_loss.detach().squeeze()
                    r1_loss = args.r1 * args.d_reg_every * r1_loss / args.num_accumulate
                    r1_loss.backward()
                d_optim.step()
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_pop()

        if args.augment and args.augment_p == 0:
            with profiler.record_function("ADA"):
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_push("ADA")
                ada_augment += th.tensor((th.sign(real_pred).sum().item(), real_pred.shape[0]), device=device)
                ada_augment = reduce_sum(ada_augment)

                if ada_augment[1] > 255:
                    pred_signs, n_pred = ada_augment.tolist()

                    r_t_stat = pred_signs / n_pred
                    loss_dict["Rt"] = th.tensor(r_t_stat, device=device).float()
                    if r_t_stat > args.ada_target:
                        sign = 1
                    else:
                        sign = -1

                    ada_aug_p += sign * ada_aug_step * n_pred
                    ada_aug_p = min(1, max(0, ada_aug_p))
                    ada_augment.mul_(0)
                    loss_dict["Augment"] = th.tensor(ada_aug_p, device=device).float()
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_pop()

        with profiler.record_function("G train"):
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("G train")
            requires_grad(generator, True)
            requires_grad(discriminator, False)

            generator.zero_grad()
            for _ in range(args.num_accumulate):
                noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)
                fake_img, _ = generator(noise)
                if args.augment:
                    fake_img, _ = augment(fake_img, ada_aug_p)
                fake_pred = discriminator(fake_img)
                g_loss = g_non_saturating_loss(fake_pred)
                loss_dict["Generator"] += g_loss.detach()
                g_loss /= args.num_accumulate
                g_loss.backward()
            g_optim.step()
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()

        if args.path_regularize > 0 and i % args.g_reg_every == 0:
            with profiler.record_function("G reg"):
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_push("G reg")
                generator.zero_grad()
                for _ in range(args.num_accumulate):
                    path_loss, mean_path_length = g_path_length_regularization(generator, mean_path_length, args)
                    loss_dict["Path Length Regularization"] += path_loss.detach()
                    path_loss = args.path_regularize * args.g_reg_every * path_loss / args.num_accumulate
                    path_loss.backward()
                g_optim.step()
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_pop()

        with profiler.record_function("Log / Eval"):
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("G accum")
            accumulate(g_ema, g_module)
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()

            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("Log / Eval")
            loss_reduced = reduce_loss_dict(loss_dict)
            log_dict = {k: v.mean().item() / args.num_accumulate for k, v in loss_reduced.items() if v != 0}
            log_dict["Tick Length"] = time.time() - tick_start

            if get_rank() == 0:
                with th.no_grad():
                    if args.log_spec_norm:
                        G_norms = []
                        for name, spec_norm in g_module.named_buffers():
                            if "spectral_norm" in name:
                                G_norms.append(spec_norm.cpu().numpy())
                        G_norms = np.array(G_norms)
                        D_norms = []
                        for name, spec_norm in d_module.named_buffers():
                            if "spectral_norm" in name:
                                D_norms.append(spec_norm.cpu().numpy())
                        D_norms = np.array(D_norms)
                        log_dict[f"Spectral Norms/G min spectral norm"] = np.log(G_norms).min()
                        log_dict[f"Spectral Norms/G mean spectral norm"] = np.log(G_norms).mean()
                        log_dict[f"Spectral Norms/G max spectral norm"] = np.log(G_norms).max()
                        log_dict[f"Spectral Norms/D min spectral norm"] = np.log(D_norms).min()
                        log_dict[f"Spectral Norms/D mean spectral norm"] = np.log(D_norms).mean()
                        log_dict[f"Spectral Norms/D max spectral norm"] = np.log(D_norms).max()

                    if args.img_every != -1 and i % args.img_every == 0:
                        g_ema.eval()
                        sample = []
                        for sub in range(0, len(sample_z), args.batch_size):
                            subsample, _ = g_ema([sample_z[sub : sub + args.batch_size]])
                            sample.append(subsample.detach().cpu())
                        sample = th.cat(sample)
                        grid = utils.make_grid(sample, nrow=10, normalize=True, range=(-1, 1))
                        log_dict["Generated Images EMA"] = [wandb.Image(grid, caption=f"Step {i}")]

                    if args.eval_every != -1 and i % args.eval_every == 0:
                        fid_dict = validation.fid(
                            g_ema, args.val_batch_size, args.fid_n_sample, args.fid_truncation, args.name
                        )

                        fid = fid_dict["FID"]
                        fids.append(fid)
                        density = fid_dict["Density"]
                        coverage = fid_dict["Coverage"]

                        ppl = validation.ppl(
                            g_ema,
                            args.val_batch_size,
                            args.ppl_n_sample,
                            args.ppl_space,
                            args.ppl_crop,
                            args.latent_size,
                        )

                        log_dict["Evaluation/FID"] = fid
                        log_dict["Sweep/FID_smooth"] = gaussian_filter(np.array(fids), [5])[-1]
                        log_dict["Evaluation/Density"] = density
                        log_dict["Evaluation/Coverage"] = coverage
                        log_dict["Evaluation/PPL"] = ppl

                        gc.collect()
                        th.cuda.empty_cache()

                    wandb.log(log_dict)

                    if args.eval_every != -1:
                        description = (
                            f"FID: {fid:.4f}   PPL: {ppl:.4f}   Dens: {density:.4f}   Cov: {coverage:.4f}   "
                            + f"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}"
                        )
                    else:
                        description = f"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}"
                    if "Augment" in log_dict:
                        description += f"   Aug: {log_dict['Augment']:.4f}"  #   Rt: {log_dict['Rt']:.4f}"
                    if "R1 Penalty" in log_dict:
                        description += f"   R1: {log_dict['R1 Penalty']:.4f}"
                    if "Path Length Regularization" in log_dict:
                        description += f"   Path: {log_dict['Path Length Regularization']:.4f}"
                    pbar.set_description(description)

                    if i % args.checkpoint_every == 0:
                        check_name = "-".join(
                            [
                                args.name,
                                args.wbname,
                                wandb.run.dir.split("/")[-1].split("-")[-1],
                                # str(int(fid)),
                                str(args.size),
                                str(i).zfill(6),
                            ]
                        )
                        th.save(
                            {
                                "g": g_module.state_dict(),
                                "d": d_module.state_dict(),
                                # "cl": cl_module.state_dict(),
                                "g_ema": g_ema.state_dict(),
                                "g_optim": g_optim.state_dict(),
                                "d_optim": d_optim.state_dict(),
                            },
                            f"/home/hans/modelzoo/maua-sg2/{check_name}.pt",
                        )

            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()  # iteration range

        gpu_profile(frame=sys._getframe(), event="line", arg=None)

    if args.nsys_iter != -1:
        th.cuda.cudart().cudaProfilerStop()


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()

    parser.add_argument("--wbname", type=str, required=True)
    parser.add_argument("--wbproj", type=str, required=True)
    parser.add_argument("--wbgroup", type=str, default=None)

    # data options
    parser.add_argument("--path", type=str, required=True)
    parser.add_argument("--vflip", type=bool, default=False)
    parser.add_argument("--hflip", type=bool, default=True)

    # training options
    parser.add_argument("--batch_size", type=int, default=12)
    parser.add_argument("--num_accumulate", type=int, default=1)

    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--transfer_mapping_only", type=bool, default=False)
    parser.add_argument("--start_iter", type=int, default=0)
    parser.add_argument("--iter", type=int, default=60_000)

    # model options
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--min_rgb_size", type=int, default=4)
    parser.add_argument("--latent_size", type=int, default=512)
    parser.add_argument("--n_mlp", type=int, default=8)
    parser.add_argument("--n_sample", type=int, default=60)
    parser.add_argument("--constant_input", type=bool, default=False)
    parser.add_argument("--channel_multiplier", type=int, default=2)
    parser.add_argument("--d_skip", type=bool, default=True)

    # optimizer options
    parser.add_argument("--lr", type=float, default=0.002)
    parser.add_argument("--d_lr_ratio", type=float, default=1.0)
    parser.add_argument("--lookahead", type=bool, default=True)
    parser.add_argument("--la_steps", type=float, default=500)
    parser.add_argument("--la_alpha", type=float, default=0.5)

    # loss options
    parser.add_argument("--r1", type=float, default=1e-5)
    parser.add_argument("--path_regularize", type=float, default=1)
    parser.add_argument("--path_batch_shrink", type=int, default=2)
    parser.add_argument("--d_reg_every", type=int, default=16)
    parser.add_argument("--g_reg_every", type=int, default=4)
    parser.add_argument("--mixing_prob", type=float, default=0.666)

    # augmentation options
    parser.add_argument("--augment", type=bool, default=False)
    parser.add_argument("--contrastive", type=float, default=0)
    parser.add_argument("--balanced_consistency", type=float, default=0)
    parser.add_argument("--augment_p", type=float, default=0)
    parser.add_argument("--ada_target", type=float, default=0.6)
    parser.add_argument("--ada_length", type=int, default=40_000)

    # validation options
    parser.add_argument("--val_batch_size", type=int, default=6)
    parser.add_argument("--fid_n_sample", type=int, default=2500)
    parser.add_argument("--fid_truncation", type=float, default=None)
    parser.add_argument("--ppl_space", choices=["z", "w"], default="w")
    parser.add_argument("--ppl_n_sample", type=int, default=1250)
    parser.add_argument("--ppl_crop", type=bool, default=False)

    # logging options
    parser.add_argument("--log_spec_norm", type=bool, default=False)
    parser.add_argument("--img_every", type=int, default=1000)
    parser.add_argument("--eval_every", type=int, default=1000)
    parser.add_argument("--checkpoint_every", type=int, default=1000)

    # (multi-)GPU options
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument("--cudnn_benchmark", type=bool, default=True)
    parser.add_argument("--nsys_iter", type=int, default=-1)
    parser.add_argument("--th_prof", action="store_true")
    parser.add_argument("--prof_gpu", action="store_true")

    args = parser.parse_args()

    with th.autograd.profiler.profile(
        enabled=args.th_prof, use_cuda=True, record_shapes=True, profile_memory=True, with_stack=True
    ) as prof:
        with profiler.record_function("init"):
            if args.balanced_consistency > 0 or args.contrastive > 0:
                args.augment = True
            args.name = os.path.splitext(os.path.basename(args.path))[0]
            args.r1 = args.r1 * args.size ** 2

            args.num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
            th.backends.cudnn.benchmark = args.cudnn_benchmark
            args.distributed = args.num_gpus > 1

            if args.distributed:
                th.cuda.set_device(args.local_rank)
                th.distributed.init_process_group(backend="nccl", init_method="env://")
                synchronize()

            generator = Generator(
                args.size,
                args.latent_size,
                args.n_mlp,
                channel_multiplier=args.channel_multiplier,
                constant_input=args.constant_input,
                min_rgb_size=args.min_rgb_size,
            ).to(device, non_blocking=True)
            discriminator = Discriminator(
                args.size, channel_multiplier=args.channel_multiplier, use_skip=args.d_skip
            ).to(device, non_blocking=True)

            if args.log_spec_norm:
                for name, parameter in generator.named_parameters():
                    if "weight" in name and parameter.squeeze().dim() > 1:
                        mod = generator
                        for attr in name.replace(".weight", "").split("."):
                            mod = getattr(mod, attr)
                        validation.track_spectral_norm(mod)
                for name, parameter in discriminator.named_parameters():
                    if "weight" in name and parameter.squeeze().dim() > 1:
                        mod = discriminator
                        for attr in name.replace(".weight", "").split("."):
                            mod = getattr(mod, attr)
                        validation.track_spectral_norm(mod)

            g_ema = Generator(
                args.size,
                args.latent_size,
                args.n_mlp,
                channel_multiplier=args.channel_multiplier,
                constant_input=args.constant_input,
                min_rgb_size=args.min_rgb_size,
            ).to(device, non_blocking=True)
            g_ema.requires_grad_(False)
            g_ema.eval()
            accumulate(g_ema, generator, 0)

            if args.contrastive > 0:
                contrast_learner = ContrastiveLearner(
                    discriminator,
                    args.size,
                    augment_fn=nn.Sequential(
                        nn.ReflectionPad2d(int((math.sqrt(2) - 1) * args.size / 4)),  # zoom out
                        augs.RandomHorizontalFlip(),
                        RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.1),
                        RandomApply(augs.RandomRotation(180), p=0.1),
                        augs.RandomResizedCrop(size=(args.size, args.size), scale=(1, 1), ratio=(1, 1)),
                        RandomApply(
                            augs.RandomResizedCrop(size=(args.size, args.size), scale=(0.5, 0.9)), p=0.1
                        ),  # zoom in
                        RandomApply(augs.RandomErasing(), p=0.1),
                    ),
                    hidden_layer=(-1, 0),
                )
            else:
                contrast_learner = None

            g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
            d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)

            g_optim = th.optim.Adam(
                generator.parameters(), lr=args.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
            )
            d_optim = th.optim.Adam(
                discriminator.parameters(),
                lr=args.lr * d_reg_ratio * args.d_lr_ratio,
                betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
            )

            if args.lookahead:
                g_optim = LookaheadMinimax(
                    g_optim, d_optim, la_steps=args.la_steps, la_alpha=args.la_alpha, accumulate=args.num_accumulate
                )

            if args.checkpoint is not None:
                print("load model:", args.checkpoint)

                checkpoint = th.load(args.checkpoint, map_location=lambda storage, loc: storage)

                try:
                    ckpt_name = os.path.basename(args.checkpoint)
                    args.start_iter = int(os.path.splitext(ckpt_name)[-1].replace(args.name, ""))
                except ValueError:
                    pass

                if args.transfer_mapping_only:
                    print("Using generator latent mapping network from checkpoint")
                    mapping_state_dict = {}
                    for key, val in checkpoint["state_dict"].items():
                        if "generator.style" in key:
                            mapping_state_dict[key.replace("generator.", "")] = val
                    generator.load_state_dict(mapping_state_dict, strict=False)
                else:
                    generator.load_state_dict(checkpoint["g"], strict=False)
                    g_ema.load_state_dict(checkpoint["g_ema"], strict=False)

                    discriminator.load_state_dict(checkpoint["d"], strict=False)

                    if args.lookahead:
                        g_optim.load_state_dict(checkpoint["g_optim"], checkpoint["d_optim"])
                    else:
                        g_optim.load_state_dict(checkpoint["g_optim"])
                        d_optim.load_state_dict(checkpoint["d_optim"])

                del checkpoint
                th.cuda.empty_cache()

            if args.distributed:
                generator = nn.parallel.DistributedDataParallel(
                    generator,
                    device_ids=[args.local_rank],
                    output_device=args.local_rank,
                    broadcast_buffers=False,
                    find_unused_parameters=True,
                )

                discriminator = nn.parallel.DistributedDataParallel(
                    discriminator,
                    device_ids=[args.local_rank],
                    output_device=args.local_rank,
                    broadcast_buffers=False,
                    find_unused_parameters=True,
                )

                if contrast_learner is not None:
                    contrast_learner = nn.parallel.DistributedDataParallel(
                        contrast_learner,
                        device_ids=[args.local_rank],
                        output_device=args.local_rank,
                        broadcast_buffers=False,
                        find_unused_parameters=True,
                    )

            transform = transforms.Compose(
                [
                    transforms.RandomVerticalFlip(p=0.5 if args.vflip else 0),
                    transforms.RandomHorizontalFlip(p=0.5 if args.hflip else 0),
                    transforms.ToTensor(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                ]
            )

            dataset = MultiResolutionDataset(args.path, transform, args.size)
            loader = data.DataLoader(
                dataset,
                batch_size=args.batch_size,
                sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),
                num_workers=0,
                drop_last=True,
                pin_memory=True,
            )

            if get_rank() == 0:
                validation.get_dataset_inception_features(loader, args.name, args.size)
                if args.wbgroup is None:
                    wandb.init(project=args.wbproj, name=args.wbname, config=vars(args))
                else:
                    wandb.init(project=args.wbproj, group=args.wbgroup, name=args.wbname, config=vars(args))

        if args.prof_gpu:
            os.environ["GPU_DEBUG"] = str(args.local_rank)
            import sys
            from gpu_profile import gpu_profile

            sys.settrace(gpu_profile)
        train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema)

    if args.th_prof:
        print(prof.total_average())
        print("cuda_memory_usage", prof.table(sort_by="cuda_memory_usage", row_limit=20))
        prof.export_chrome_trace(f"{args.name}_gpu{args.local_rank}.trace")