run_train.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved


from math import dist
import sys
import os
import click
import re
import json
import glob
import tempfile
import torch
import dnnlib
import hydra

from datetime import date
from training import training_loop
from metrics import metric_main
from torch_utils import training_stats, custom_ops, distributed_utils
from torch_utils.distributed_utils import get_init_file, get_shared_folder
from omegaconf import DictConfig, OmegaConf

#----------------------------------------------------------------------------

class UserError(Exception):
    pass

#----------------------------------------------------------------------------

def setup_training_loop_kwargs(cfg):
    args = OmegaConf.create({})

    # ------------------------------------------
    # General options: gpus, snap, metrics, seed
    # ------------------------------------------
    args.rank       = 0
    args.gpu        = 0
    args.num_gpus   = torch.cuda.device_count() if cfg.gpus is None else cfg.gpus
    args.nodes      = cfg.nodes if cfg.nodes is not None else 1
    args.world_size = 1
    
    args.dist_url   = 'env://'
    args.launcher   = cfg.launcher
    args.partition  = cfg.partition
    args.comment    = cfg.comment
    args.timeout    = 4320 if cfg.timeout is None else cfg.timeout
    args.job_dir    = ''

    if cfg.snap is None:
        cfg.snap = 50
    assert isinstance(cfg.snap, int)
    if cfg.snap < 1:
        raise UserError('snap must be at least 1')
    args.image_snapshot_ticks = cfg.imgsnap
    args.network_snapshot_ticks = cfg.snap
    if hasattr(cfg, 'ucp'):
        args.update_cam_prior_ticks = cfg.ucp

    if cfg.metrics is None:
        cfg.metrics = ['fid50k_full']
    cfg.metrics = list(cfg.metrics)
    if not all(metric_main.is_valid_metric(metric) for metric in cfg.metrics):
        raise UserError('\n'.join(['metrics can only contain the following values:'] + metric_main.list_valid_metrics()))
    args.metrics = cfg.metrics

    if cfg.seed is None:
        cfg.seed = 0
    assert isinstance(cfg.seed, int)
    args.random_seed = cfg.seed

    # -----------------------------------
    # Dataset: data, cond, subset, mirror
    # ----------------------------------- 

    assert cfg.data is not None
    assert isinstance(cfg.data, str)
    args.update({"training_set_kwargs": dict(class_name='training.dataset.ImageFolderDataset', path=cfg.data, resolution=cfg.resolution, use_labels=True, max_size=None, xflip=False)})
    args.update({"data_loader_kwargs": dict(pin_memory=True, num_workers=3, prefetch_factor=2)})
    args.generation_with_image = getattr(cfg, 'generate_with_image', False)
    try:
        training_set = dnnlib.util.construct_class_by_name(**args.training_set_kwargs) # subclass of training.dataset.Dataset
        args.training_set_kwargs.resolution = training_set.resolution                  # be explicit about resolution
        args.training_set_kwargs.use_labels = training_set.has_labels                  # be explicit about labels
        args.training_set_kwargs.max_size = len(training_set)                          # be explicit about dataset size
        desc = training_set.name
        del training_set # conserve memory
    except IOError as err:
        raise UserError(f'data: {err}')

    if cfg.cond is None:
        cfg.cond = False
    assert isinstance(cfg.cond, bool)
    if cfg.cond:
        if not args.training_set_kwargs.use_labels:
            raise UserError('cond=True requires labels specified in dataset.json')
        desc += '-cond'
    else:
        args.training_set_kwargs.use_labels = False

    if cfg.subset is not None:
        assert isinstance(cfg.subset, int)
        if not 1 <= cfg.subset <= args.training_set_kwargs.max_size:
            raise UserError(f'subset must be between 1 and {args.training_set_kwargs.max_size}')
        desc += f'-subset{cfg.subset}'
        if cfg.subset < args.training_set_kwargs.max_size:
            args.training_set_kwargs.max_size = cfg.subset
            args.training_set_kwargs.random_seed = args.random_seed

    if cfg.mirror is None:
        cfg.mirror = False
    assert isinstance(cfg.mirror, bool)
    if cfg.mirror:
        desc += '-mirror'
        args.training_set_kwargs.xflip = True

    # ------------------------------------
    # Base config: cfg, model, gamma, kimg, batch
    # ------------------------------------
    if cfg.auto:
        cfg.spec.name = 'auto'
    desc += f'-{cfg.spec.name}'
    desc += f'-{cfg.model.name}'
    if cfg.spec.name == 'auto':
        res = args.training_set_kwargs.resolution
        cfg.spec.fmaps = 1 if res >= 512 else 0.5
        cfg.spec.lrate = 0.002 if res >= 1024 else 0.0025
        cfg.spec.gamma = 0.0002 * (res ** 2) / cfg.spec.mb # heuristic formula
        cfg.spec.ema = cfg.spec.mb * 10 / 32
    
    if getattr(cfg.spec, 'lrate_disc', None) is None:
        cfg.spec.lrate_disc = cfg.spec.lrate   # use the same learning rate for discriminator

    # model (generator, discriminator)
    args.update({"G_kwargs": dict(**cfg.model.G_kwargs)})
    args.update({"D_kwargs": dict(**cfg.model.D_kwargs)})
    args.update({"G_opt_kwargs": dict(class_name='torch.optim.Adam', lr=cfg.spec.lrate, betas=[0,0.99], eps=1e-8)})
    args.update({"D_opt_kwargs": dict(class_name='torch.optim.Adam', lr=cfg.spec.lrate_disc, betas=[0,0.99], eps=1e-8)})
    args.update({"loss_kwargs": dict(class_name='training.loss.StyleGAN2Loss', r1_gamma=cfg.spec.gamma, **cfg.model.loss_kwargs)})
    
    if cfg.spec.name == 'cifar':
        args.loss_kwargs.pl_weight = 0 # disable path length regularization
        args.loss_kwargs.style_mixing_prob = 0 # disable style mixing
        args.D_kwargs.architecture = 'orig' # disable residual skip connections

    # kimg data config
    args.spec = cfg.spec  # just keep the dict.
    args.total_kimg = cfg.spec.kimg
    args.batch_size = cfg.spec.mb
    args.batch_gpu = cfg.spec.mbstd
    args.ema_kimg = cfg.spec.ema
    args.ema_rampup = cfg.spec.ramp
    
    # ---------------------------------------------------
    # Discriminator augmentation: aug, p, target, augpipe
    # ---------------------------------------------------
    if cfg.aug is None:
        cfg.aug = 'ada'
    else:
        assert isinstance(cfg.aug, str)
        desc += f'-{cfg.aug}'

    if cfg.aug == 'ada':
        args.ada_target = 0.6
    elif cfg.aug == 'noaug':
        pass
    elif cfg.aug == 'fixed':
        if cfg.p is None:
            raise UserError(f'--aug={cfg.aug} requires specifying --p')
    else:
        raise UserError(f'--aug={cfg.aug} not supported')

    if cfg.p is not None:
        assert isinstance(cfg.p, float)
        if cfg.aug != 'fixed':
            raise UserError('--p can only be specified with --aug=fixed')
        if not 0 <= cfg.p <= 1:
            raise UserError('--p must be between 0 and 1')
        desc += f'-p{cfg.p:g}'
        args.augment_p = cfg.p

    if cfg.target is not None:
        assert isinstance(cfg.target, float)
        if cfg.aug != 'ada':
            raise UserError('--target can only be specified with --aug=ada')
        if not 0 <= cfg.target <= 1:
            raise UserError('--target must be between 0 and 1')
        desc += f'-target{cfg.target:g}'
        args.ada_target = cfg.target

    assert cfg.augpipe is None or isinstance(cfg.augpipe, str)
    if cfg.augpipe is None:
        cfg.augpipe = 'bgc'
    else:
        if cfg.aug == 'noaug':
            raise UserError('--augpipe cannot be specified with --aug=noaug')
        desc += f'-{cfg.augpipe}'

    augpipe_specs = {
        'blit':   dict(xflip=1, rotate90=1, xint=1),
        'geom':   dict(scale=1, rotate=1, aniso=1, xfrac=1),
        'color':  dict(brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),
        'filter': dict(imgfilter=1),
        'noise':  dict(noise=1),
        'cutout': dict(cutout=1),
        'bgc0':   dict(xint=1, scale=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),
        'bg':     dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1),
        'bgc':    dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),
        'bgcf':   dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1, imgfilter=1),
        'bgcfn':  dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1, imgfilter=1, noise=1),
        'bgcfnc': dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1, imgfilter=1, noise=1, cutout=1),
    }
    assert cfg.augpipe in augpipe_specs
    if cfg.aug != 'noaug':
        args.update({"augment_kwargs": dict(class_name='training.augment.AugmentPipe', **augpipe_specs[cfg.augpipe])})

    # ----------------------------------
    # Transfer learning: resume, freezed
    # ----------------------------------

    resume_specs = {
        'ffhq256':     'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res256-mirror-paper256-noaug.pkl',
        'ffhq512':     'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res512-mirror-stylegan2-noaug.pkl',
        'ffhq1024':    'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res1024-mirror-stylegan2-noaug.pkl',
        'celebahq256': 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/celebahq-res256-mirror-paper256-kimg100000-ada-target0.5.pkl',
        'lsundog256':  'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/lsundog-res256-paper256-kimg100000-noaug.pkl',
    }

    assert cfg.resume is None or isinstance(cfg.resume, str)
    if cfg.resume is None:
        cfg.resume = 'noresume'
    elif cfg.resume == 'noresume':
        desc += '-noresume'
    elif cfg.resume in resume_specs:
        desc += f'-resume{cfg.resume}'
        args.resume_pkl = resume_specs[cfg.resume] # predefined url
    else:
        desc += '-resumecustom'
        args.resume_pkl = cfg.resume # custom path or url

    if cfg.resume != 'noresume':
        args.ada_kimg = 100 # make ADA react faster at the beginning
        args.ema_rampup = None # disable EMA rampup

    if cfg.freezed is not None:
        assert isinstance(cfg.freezed, int)
        if not cfg.freezed >= 0:
            raise UserError('--freezed must be non-negative')
        desc += f'-freezed{cfg.freezed:d}'
        args.D_kwargs.block_kwargs.freeze_layers = cfg.freezed

    # -------------------------------------------------
    # Performance options: fp32, nhwc, nobench, workers
    # -------------------------------------------------
    args.num_fp16_res = cfg.num_fp16_res
    if cfg.fp32 is None:
        cfg.fp32 = False
    assert isinstance(cfg.fp32, bool)
    if cfg.fp32:
        args.G_kwargs.synthesis_kwargs.num_fp16_res = args.D_kwargs.num_fp16_res = 0
        args.G_kwargs.synthesis_kwargs.conv_clamp = args.D_kwargs.conv_clamp = None

    if cfg.nhwc is None:
        cfg.nhwc = False
    assert isinstance(cfg.nhwc, bool)
    if cfg.nhwc:
        args.G_kwargs.synthesis_kwargs.fp16_channels_last = args.D_kwargs.block_kwargs.fp16_channels_last = True

    if cfg.nobench is None:
        cfg.nobench = False
    assert isinstance(cfg.nobench, bool)
    if cfg.nobench:
        args.cudnn_benchmark = False

    if cfg.allow_tf32 is None:
        cfg.allow_tf32 = False
    assert isinstance(cfg.allow_tf32, bool)
    args.allow_tf32 = cfg.allow_tf32

    if cfg.workers is not None:
        assert isinstance(cfg.workers, int)
        if not cfg.workers >= 1:
            raise UserError('--workers must be at least 1')
        args.data_loader_kwargs.num_workers = cfg.workers

    args.debug = cfg.debug
    if getattr(cfg, "prefix", None) is not None:
        desc = cfg.prefix + '-' + desc
    return desc, args

#----------------------------------------------------------------------------

def subprocess_fn(rank, args):
    if not args.debug:
        dnnlib.util.Logger(file_name=os.path.join(args.run_dir, 'log.txt'), file_mode='a', should_flush=True)

    # Init torch.distributed.
    distributed_utils.init_distributed_mode(rank, args)
    if args.rank != 0:
        custom_ops.verbosity = 'none'
    
    # Execute training loop.
    training_loop.training_loop(**args)
    
#----------------------------------------------------------------------------

class CommaSeparatedList(click.ParamType):
    name = 'list'

    def convert(self, value, param, ctx):
        _ = param, ctx
        if value is None or value.lower() == 'none' or value == '':
            return []
        return value.split(',')


@hydra.main(config_path="conf", config_name="config")
def main(cfg: DictConfig):
    
    outdir = cfg.outdir

    # Setup training options
    run_desc, args = setup_training_loop_kwargs(cfg)
    
    # Pick output directory.
    prev_run_dirs = []
    if os.path.isdir(outdir):
        prev_run_dirs = [x for x in os.listdir(outdir) if os.path.isdir(os.path.join(outdir, x))]
    
    if cfg.resume_run is None:
        prev_run_ids = [re.match(r'^\d+', x) for x in prev_run_dirs]
        prev_run_ids = [int(x.group()) for x in prev_run_ids if x is not None]
        cur_run_id = max(prev_run_ids, default=-1) + 1
    else:
        cur_run_id = cfg.resume_run
        
    args.run_dir = os.path.join(outdir, f'{cur_run_id:05d}-{run_desc}')
    print(outdir, args.run_dir)

    if cfg.resume_run is not None:
        pkls = sorted(glob.glob(args.run_dir + '/network*.pkl'))
        if len(pkls) > 0:
            args.resume_pkl = pkls[-1]
            args.resume_start = int(args.resume_pkl.split('-')[-1][:-4]) * 1000
        else:
            args.resume_start = 0

    # Print options.
    print()
    print('Training options:')
    print(OmegaConf.to_yaml(args))
    print()
    print(f'Output directory:   {args.run_dir}')
    print(f'Training data:      {args.training_set_kwargs.path}')
    print(f'Training duration:  {args.total_kimg} kimg')
    print(f'Number of images:   {args.training_set_kwargs.max_size}')
    print(f'Image resolution:   {args.training_set_kwargs.resolution}')
    print(f'Conditional model:  {args.training_set_kwargs.use_labels}')
    print(f'Dataset x-flips:    {args.training_set_kwargs.xflip}')
    print()

    # Dry run?
    if cfg.dry_run:
        print('Dry run; exiting.')
        return

    # Create output directory.
    print('Creating output directory...')
    if not os.path.exists(args.run_dir):
        os.makedirs(args.run_dir)
        with open(os.path.join(args.run_dir, 'training_options.yaml'), 'wt') as fp:
            OmegaConf.save(config=args, f=fp.name)

    # Launch processes.    
    print('Launching processes...')
    if (args.launcher == 'spawn') and (args.num_gpus > 1):
        args.dist_url = distributed_utils.get_init_file().as_uri()
        torch.multiprocessing.set_start_method('spawn')
        torch.multiprocessing.spawn(fn=subprocess_fn, args=(args,), nprocs=args.num_gpus)
    else:
        subprocess_fn(rank=0, args=args)

#----------------------------------------------------------------------------

if __name__ == "__main__":
    if os.getenv('SLURM_ARGS') is not None:
        # deparcated launcher for slurm jobs.
        slurm_arg = eval(os.getenv('SLURM_ARGS'))
        all_args = sys.argv[1:]
        print(slurm_arg)
        print(all_args)

        from launcher import launch
        launch(slurm_arg, all_args)
    
    else:
        main() # pylint: disable=no-value-for-parameter

#----------------------------------------------------------------------------