facebook/StyleNeRF
34
1# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved2 3 4from math import dist5import sys6import os7import click8import re9import json10import glob11import tempfile12import torch13import dnnlib14import hydra15 16from datetime import date17from training import training_loop18from metrics import metric_main19from torch_utils import training_stats, custom_ops, distributed_utils20from torch_utils.distributed_utils import get_init_file, get_shared_folder21from omegaconf import DictConfig, OmegaConf22 23#----------------------------------------------------------------------------24 25class UserError(Exception):26 pass27 28#----------------------------------------------------------------------------29 30def setup_training_loop_kwargs(cfg):31 args = OmegaConf.create({})32 33 # ------------------------------------------34 # General options: gpus, snap, metrics, seed35 # ------------------------------------------36 args.rank = 037 args.gpu = 038 args.num_gpus = torch.cuda.device_count() if cfg.gpus is None else cfg.gpus39 args.nodes = cfg.nodes if cfg.nodes is not None else 140 args.world_size = 141 42 args.dist_url = 'env://'43 args.launcher = cfg.launcher44 args.partition = cfg.partition45 args.comment = cfg.comment46 args.timeout = 4320 if cfg.timeout is None else cfg.timeout47 args.job_dir = ''48 49 if cfg.snap is None:50 cfg.snap = 5051 assert isinstance(cfg.snap, int)52 if cfg.snap < 1:53 raise UserError('snap must be at least 1')54 args.image_snapshot_ticks = cfg.imgsnap55 args.network_snapshot_ticks = cfg.snap56 if hasattr(cfg, 'ucp'):57 args.update_cam_prior_ticks = cfg.ucp58 59 if cfg.metrics is None:60 cfg.metrics = ['fid50k_full']61 cfg.metrics = list(cfg.metrics)62 if not all(metric_main.is_valid_metric(metric) for metric in cfg.metrics):63 raise UserError('\n'.join(['metrics can only contain the following values:'] + metric_main.list_valid_metrics()))64 args.metrics = cfg.metrics65 66 if cfg.seed is None:67 cfg.seed = 068 assert isinstance(cfg.seed, int)69 args.random_seed = cfg.seed70 71 # -----------------------------------72 # Dataset: data, cond, subset, mirror73 # ----------------------------------- 74 75 assert cfg.data is not None76 assert isinstance(cfg.data, str)77 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)})78 args.update({"data_loader_kwargs": dict(pin_memory=True, num_workers=3, prefetch_factor=2)})79 args.generation_with_image = getattr(cfg, 'generate_with_image', False)80 try:81 training_set = dnnlib.util.construct_class_by_name(**args.training_set_kwargs) # subclass of training.dataset.Dataset82 args.training_set_kwargs.resolution = training_set.resolution # be explicit about resolution83 args.training_set_kwargs.use_labels = training_set.has_labels # be explicit about labels84 args.training_set_kwargs.max_size = len(training_set) # be explicit about dataset size85 desc = training_set.name86 del training_set # conserve memory87 except IOError as err:88 raise UserError(f'data: {err}')89 90 if cfg.cond is None:91 cfg.cond = False92 assert isinstance(cfg.cond, bool)93 if cfg.cond:94 if not args.training_set_kwargs.use_labels:95 raise UserError('cond=True requires labels specified in dataset.json')96 desc += '-cond'97 else:98 args.training_set_kwargs.use_labels = False99 100 if cfg.subset is not None:101 assert isinstance(cfg.subset, int)102 if not 1 <= cfg.subset <= args.training_set_kwargs.max_size:103 raise UserError(f'subset must be between 1 and {args.training_set_kwargs.max_size}')104 desc += f'-subset{cfg.subset}'105 if cfg.subset < args.training_set_kwargs.max_size:106 args.training_set_kwargs.max_size = cfg.subset107 args.training_set_kwargs.random_seed = args.random_seed108 109 if cfg.mirror is None:110 cfg.mirror = False111 assert isinstance(cfg.mirror, bool)112 if cfg.mirror:113 desc += '-mirror'114 args.training_set_kwargs.xflip = True115 116 # ------------------------------------117 # Base config: cfg, model, gamma, kimg, batch118 # ------------------------------------119 if cfg.auto:120 cfg.spec.name = 'auto'121 desc += f'-{cfg.spec.name}'122 desc += f'-{cfg.model.name}'123 if cfg.spec.name == 'auto':124 res = args.training_set_kwargs.resolution125 cfg.spec.fmaps = 1 if res >= 512 else 0.5126 cfg.spec.lrate = 0.002 if res >= 1024 else 0.0025127 cfg.spec.gamma = 0.0002 * (res ** 2) / cfg.spec.mb # heuristic formula128 cfg.spec.ema = cfg.spec.mb * 10 / 32129 130 if getattr(cfg.spec, 'lrate_disc', None) is None:131 cfg.spec.lrate_disc = cfg.spec.lrate # use the same learning rate for discriminator132 133 # model (generator, discriminator)134 args.update({"G_kwargs": dict(**cfg.model.G_kwargs)})135 args.update({"D_kwargs": dict(**cfg.model.D_kwargs)})136 args.update({"G_opt_kwargs": dict(class_name='torch.optim.Adam', lr=cfg.spec.lrate, betas=[0,0.99], eps=1e-8)})137 args.update({"D_opt_kwargs": dict(class_name='torch.optim.Adam', lr=cfg.spec.lrate_disc, betas=[0,0.99], eps=1e-8)})138 args.update({"loss_kwargs": dict(class_name='training.loss.StyleGAN2Loss', r1_gamma=cfg.spec.gamma, **cfg.model.loss_kwargs)})139 140 if cfg.spec.name == 'cifar':141 args.loss_kwargs.pl_weight = 0 # disable path length regularization142 args.loss_kwargs.style_mixing_prob = 0 # disable style mixing143 args.D_kwargs.architecture = 'orig' # disable residual skip connections144 145 # kimg data config146 args.spec = cfg.spec # just keep the dict.147 args.total_kimg = cfg.spec.kimg148 args.batch_size = cfg.spec.mb149 args.batch_gpu = cfg.spec.mbstd150 args.ema_kimg = cfg.spec.ema151 args.ema_rampup = cfg.spec.ramp152 153 # ---------------------------------------------------154 # Discriminator augmentation: aug, p, target, augpipe155 # ---------------------------------------------------156 if cfg.aug is None:157 cfg.aug = 'ada'158 else:159 assert isinstance(cfg.aug, str)160 desc += f'-{cfg.aug}'161 162 if cfg.aug == 'ada':163 args.ada_target = 0.6164 elif cfg.aug == 'noaug':165 pass166 elif cfg.aug == 'fixed':167 if cfg.p is None:168 raise UserError(f'--aug={cfg.aug} requires specifying --p')169 else:170 raise UserError(f'--aug={cfg.aug} not supported')171 172 if cfg.p is not None:173 assert isinstance(cfg.p, float)174 if cfg.aug != 'fixed':175 raise UserError('--p can only be specified with --aug=fixed')176 if not 0 <= cfg.p <= 1:177 raise UserError('--p must be between 0 and 1')178 desc += f'-p{cfg.p:g}'179 args.augment_p = cfg.p180 181 if cfg.target is not None:182 assert isinstance(cfg.target, float)183 if cfg.aug != 'ada':184 raise UserError('--target can only be specified with --aug=ada')185 if not 0 <= cfg.target <= 1:186 raise UserError('--target must be between 0 and 1')187 desc += f'-target{cfg.target:g}'188 args.ada_target = cfg.target189 190 assert cfg.augpipe is None or isinstance(cfg.augpipe, str)191 if cfg.augpipe is None:192 cfg.augpipe = 'bgc'193 else:194 if cfg.aug == 'noaug':195 raise UserError('--augpipe cannot be specified with --aug=noaug')196 desc += f'-{cfg.augpipe}'197 198 augpipe_specs = {199 'blit': dict(xflip=1, rotate90=1, xint=1),200 'geom': dict(scale=1, rotate=1, aniso=1, xfrac=1),201 'color': dict(brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),202 'filter': dict(imgfilter=1),203 'noise': dict(noise=1),204 'cutout': dict(cutout=1),205 'bgc0': dict(xint=1, scale=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),206 'bg': dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1),207 '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),208 '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),209 '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),210 '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),211 }212 assert cfg.augpipe in augpipe_specs213 if cfg.aug != 'noaug':214 args.update({"augment_kwargs": dict(class_name='training.augment.AugmentPipe', **augpipe_specs[cfg.augpipe])})215 216 # ----------------------------------217 # Transfer learning: resume, freezed218 # ----------------------------------219 220 resume_specs = {221 'ffhq256': 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res256-mirror-paper256-noaug.pkl',222 'ffhq512': 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res512-mirror-stylegan2-noaug.pkl',223 'ffhq1024': 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res1024-mirror-stylegan2-noaug.pkl',224 'celebahq256': 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/celebahq-res256-mirror-paper256-kimg100000-ada-target0.5.pkl',225 'lsundog256': 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/lsundog-res256-paper256-kimg100000-noaug.pkl',226 }227 228 assert cfg.resume is None or isinstance(cfg.resume, str)229 if cfg.resume is None:230 cfg.resume = 'noresume'231 elif cfg.resume == 'noresume':232 desc += '-noresume'233 elif cfg.resume in resume_specs:234 desc += f'-resume{cfg.resume}'235 args.resume_pkl = resume_specs[cfg.resume] # predefined url236 else:237 desc += '-resumecustom'238 args.resume_pkl = cfg.resume # custom path or url239 240 if cfg.resume != 'noresume':241 args.ada_kimg = 100 # make ADA react faster at the beginning242 args.ema_rampup = None # disable EMA rampup243 244 if cfg.freezed is not None:245 assert isinstance(cfg.freezed, int)246 if not cfg.freezed >= 0:247 raise UserError('--freezed must be non-negative')248 desc += f'-freezed{cfg.freezed:d}'249 args.D_kwargs.block_kwargs.freeze_layers = cfg.freezed250 251 # -------------------------------------------------252 # Performance options: fp32, nhwc, nobench, workers253 # -------------------------------------------------254 args.num_fp16_res = cfg.num_fp16_res255 if cfg.fp32 is None:256 cfg.fp32 = False257 assert isinstance(cfg.fp32, bool)258 if cfg.fp32:259 args.G_kwargs.synthesis_kwargs.num_fp16_res = args.D_kwargs.num_fp16_res = 0260 args.G_kwargs.synthesis_kwargs.conv_clamp = args.D_kwargs.conv_clamp = None261 262 if cfg.nhwc is None:263 cfg.nhwc = False264 assert isinstance(cfg.nhwc, bool)265 if cfg.nhwc:266 args.G_kwargs.synthesis_kwargs.fp16_channels_last = args.D_kwargs.block_kwargs.fp16_channels_last = True267 268 if cfg.nobench is None:269 cfg.nobench = False270 assert isinstance(cfg.nobench, bool)271 if cfg.nobench:272 args.cudnn_benchmark = False273 274 if cfg.allow_tf32 is None:275 cfg.allow_tf32 = False276 assert isinstance(cfg.allow_tf32, bool)277 args.allow_tf32 = cfg.allow_tf32278 279 if cfg.workers is not None:280 assert isinstance(cfg.workers, int)281 if not cfg.workers >= 1:282 raise UserError('--workers must be at least 1')283 args.data_loader_kwargs.num_workers = cfg.workers284 285 args.debug = cfg.debug286 if getattr(cfg, "prefix", None) is not None:287 desc = cfg.prefix + '-' + desc288 return desc, args289 290#----------------------------------------------------------------------------291 292def subprocess_fn(rank, args):293 if not args.debug:294 dnnlib.util.Logger(file_name=os.path.join(args.run_dir, 'log.txt'), file_mode='a', should_flush=True)295 296 # Init torch.distributed.297 distributed_utils.init_distributed_mode(rank, args)298 if args.rank != 0:299 custom_ops.verbosity = 'none'300 301 # Execute training loop.302 training_loop.training_loop(**args)303 304#----------------------------------------------------------------------------305 306class CommaSeparatedList(click.ParamType):307 name = 'list'308 309 def convert(self, value, param, ctx):310 _ = param, ctx311 if value is None or value.lower() == 'none' or value == '':312 return []313 return value.split(',')314 315 316@hydra.main(config_path="conf", config_name="config")317def main(cfg: DictConfig):318 319 outdir = cfg.outdir320 321 # Setup training options322 run_desc, args = setup_training_loop_kwargs(cfg)323 324 # Pick output directory.325 prev_run_dirs = []326 if os.path.isdir(outdir):327 prev_run_dirs = [x for x in os.listdir(outdir) if os.path.isdir(os.path.join(outdir, x))]328 329 if cfg.resume_run is None:330 prev_run_ids = [re.match(r'^\d+', x) for x in prev_run_dirs]331 prev_run_ids = [int(x.group()) for x in prev_run_ids if x is not None]332 cur_run_id = max(prev_run_ids, default=-1) + 1333 else:334 cur_run_id = cfg.resume_run335 336 args.run_dir = os.path.join(outdir, f'{cur_run_id:05d}-{run_desc}')337 print(outdir, args.run_dir)338 339 if cfg.resume_run is not None:340 pkls = sorted(glob.glob(args.run_dir + '/network*.pkl'))341 if len(pkls) > 0:342 args.resume_pkl = pkls[-1]343 args.resume_start = int(args.resume_pkl.split('-')[-1][:-4]) * 1000344 else:345 args.resume_start = 0346 347 # Print options.348 print()349 print('Training options:')350 print(OmegaConf.to_yaml(args))351 print()352 print(f'Output directory: {args.run_dir}')353 print(f'Training data: {args.training_set_kwargs.path}')354 print(f'Training duration: {args.total_kimg} kimg')355 print(f'Number of images: {args.training_set_kwargs.max_size}')356 print(f'Image resolution: {args.training_set_kwargs.resolution}')357 print(f'Conditional model: {args.training_set_kwargs.use_labels}')358 print(f'Dataset x-flips: {args.training_set_kwargs.xflip}')359 print()360 361 # Dry run?362 if cfg.dry_run:363 print('Dry run; exiting.')364 return365 366 # Create output directory.367 print('Creating output directory...')368 if not os.path.exists(args.run_dir):369 os.makedirs(args.run_dir)370 with open(os.path.join(args.run_dir, 'training_options.yaml'), 'wt') as fp:371 OmegaConf.save(config=args, f=fp.name)372 373 # Launch processes. 374 print('Launching processes...')375 if (args.launcher == 'spawn') and (args.num_gpus > 1):376 args.dist_url = distributed_utils.get_init_file().as_uri()377 torch.multiprocessing.set_start_method('spawn')378 torch.multiprocessing.spawn(fn=subprocess_fn, args=(args,), nprocs=args.num_gpus)379 else:380 subprocess_fn(rank=0, args=args)381 382#----------------------------------------------------------------------------383 384if __name__ == "__main__":385 if os.getenv('SLURM_ARGS') is not None:386 # deparcated launcher for slurm jobs.387 slurm_arg = eval(os.getenv('SLURM_ARGS'))388 all_args = sys.argv[1:]389 print(slurm_arg)390 print(all_args)391 392 from launcher import launch393 launch(slurm_arg, all_args)394 395 else:396 main() # pylint: disable=no-value-for-parameter397 398#----------------------------------------------------------------------------399