def run_auto(dataset, data_dir, result_dir, config_id, num_gpus, resolution, total_kimg, gamma, mirror_augment, metrics, train_auto): train = EasyDict(run_func_name='training.training_loop.training_auto_loop' ) # Options for training loop. Enc = EasyDict(func_name='training.networks_stylegan2.Encoder' ) # Options for encoder network. Dec = EasyDict(func_name='training.networks_stylegan2.Decoder' ) # Options for decoder network. opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for autoencoder optimizer. loss = EasyDict( func_name='training.loss.auto_l1') # Options for autoencoder loss. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). train.data_dir = data_dir train.total_kimg = total_kimg train.image_snapshot_ticks = 10 train.network_snapshot_ticks = 125 sched.lrate = 0.003 sched.minibatch_size = 64 sched.minibatch_gpu = 64 desc = 'stylegan2-hrae' desc += '-' + dataset dataset_args = EasyDict(tfrecord_dir=dataset) dataset_args.resolution = resolution dataset_args.num_threads = 4 assert num_gpus in [1, 2, 4, 8] sc.num_gpus = num_gpus desc += '-%dgpu' % num_gpus desc += '-auto' sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(Enc_args=Enc, Dec_args=Dec, opt_args=opt, loss_args=loss) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, tf_config=tf_config) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics): train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop. G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg' ) # Options for generator loss. D_loss = EasyDict(func_name='training.loss.D_logistic_r1' ) # Options for discriminator loss. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='8k', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). try: pkl, kimg = misc.locate_latest_pkl(result_dir) train.resume_pkl = pkl train.resume_kimg = kimg except: print('Couldn\'t find valid snapshot, starting over') train.data_dir = data_dir train.total_kimg = total_kimg train.mirror_augment = mirror_augment train.image_snapshot_ticks = 1 train.network_snapshot_ticks = 2 sched.G_lrate_base = sched.D_lrate_base = 0.002 sched.minibatch_size_base = 32 sched.minibatch_gpu_base = 4 D_loss.gamma = 10 metrics = [metric_defaults[x] for x in metrics] desc = 'stylegan2' desc += '-' + dataset dataset_args = EasyDict(tfrecord_dir=dataset) assert num_gpus in [1, 2, 4, 8] sc.num_gpus = num_gpus desc += '-%dgpu' % num_gpus assert config_id in _valid_configs desc += '-' + config_id # Configs A-E: Shrink networks to match original StyleGAN. if config_id != 'config-f': G.fmap_base = D.fmap_base = 8 << 10 # Config E: Set gamma to 100 and override G & D architecture. if config_id.startswith('config-e'): D_loss.gamma = 100 if 'Gorig' in config_id: G.architecture = 'orig' if 'Gskip' in config_id: G.architecture = 'skip' # (default) if 'Gresnet' in config_id: G.architecture = 'resnet' if 'Dorig' in config_id: D.architecture = 'orig' if 'Dskip' in config_id: D.architecture = 'skip' if 'Dresnet' in config_id: D.architecture = 'resnet' # (default) # Configs A-D: Enable progressive growing and switch to networks that support it. if config_id in ['config-a', 'config-b', 'config-c', 'config-d']: sched.lod_initial_resolution = 8 sched.G_lrate_base = sched.D_lrate_base = 0.001 sched.G_lrate_dict = sched.D_lrate_dict = { 128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003 } sched.minibatch_size_base = 32 # (default) sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32} sched.minibatch_gpu_base = 4 # (default) sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4} G.synthesis_func = 'G_synthesis_stylegan_revised' D.func_name = 'training.networks_stylegan2.D_stylegan' # Configs A-C: Disable path length regularization. if config_id in ['config-a', 'config-b', 'config-c']: G_loss = EasyDict(func_name='training.loss.G_logistic_ns') # Configs A-B: Disable lazy regularization. if config_id in ['config-a', 'config-b']: train.lazy_regularization = False # Config A: Switch to original StyleGAN networks. if config_id == 'config-a': G = EasyDict(func_name='training.networks_stylegan.G_style') D = EasyDict(func_name='training.networks_stylegan.D_basic') if gamma is not None: D_loss.gamma = gamma sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def run(dataset, resolution, result_dir, DiffAugment, num_gpus, batch_size, total_kimg, ema_kimg, num_samples, gamma, fmap_base, fmap_max, latent_size, mirror_augment, impl, metrics, resume, resume_kimg, num_repeats, eval): train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop. G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. loss_args = EasyDict( func_name='training.loss.ns_r1_DiffAugment') # Options for loss. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='4k', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). train.image_snapshot_ticks = 1 train.network_snapshot_ticks = 4 sched.minibatch_gpu_base = 8 # preprocess dataset into tfrecords if necessary dataset = dataset_tool.create_dataset(dataset, resolution) train.total_kimg = total_kimg train.mirror_augment = mirror_augment metrics = [metric_defaults[x] for x in metrics] metric_args = EasyDict(cache_dir=dataset, num_repeats=num_repeats) desc = 'DiffAugment-stylegan2' if DiffAugment else 'stylegan2' dataset_args = EasyDict(tfrecord_dir=dataset, resolution=resolution, from_tfrecords=True) desc += '-' + os.path.basename(dataset) if resolution is not None: desc += '-{}'.format(resolution) if num_samples is not None: dataset_args.num_samples = num_samples desc += '-{}samples'.format(num_samples) if batch_size is not None: desc += '-batch{}'.format(batch_size) else: batch_size = 32 assert batch_size % num_gpus == 0 sc.num_gpus = num_gpus desc += '-%dgpu' % num_gpus sched.minibatch_size_base = batch_size sched.minibatch_gpu_base = batch_size // num_gpus G.impl = D.impl = impl if fmap_base is not None: G.fmap_base = D.fmap_base = fmap_base desc += '-fmap{}'.format(fmap_base) if fmap_max is not None: G.fmap_max = D.fmap_max = fmap_max desc += '-fmax{}'.format(fmap_max) if latent_size is not None: G.latent_size = G.mapping_fmaps = G.dlatent_size = latent_size desc += '-latent{}'.format(latent_size) if gamma is not None: loss_args.gamma = gamma desc += '-gamma{}'.format(gamma) if DiffAugment: loss_args.policy = DiffAugment desc += '-' + DiffAugment.replace(',', '-') sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, loss_args=loss_args) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config) kwargs.update(resume_pkl=resume, resume_kimg=resume_kimg, resume_with_new_nets=True) kwargs.update(metric_args=metric_args) if ema_kimg is not None: kwargs.update(G_ema_kimg=ema_kimg) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def run(dataset, train_dir, config, d_aug, diffaug_policy, cond, ops, jpg_data, mirror, mirror_v, \ lod_step_kimg, batch_size, resume, resume_kimg, finetune, num_gpus, ema_kimg, gamma, freezeD): # dataset (tfrecords) - preprocess or get tfr_files = file_list(os.path.dirname(dataset), 'tfr') tfr_files = [f for f in tfr_files if basename(dataset) in f] if len(tfr_files) == 0: tfr_file, total_samples = create_from_images(dataset, jpg=jpg_data) else: tfr_file = tfr_files[0] dataset_args = EasyDict(tfrecord=tfr_file, jpg_data=jpg_data) desc = basename(tfr_file).split('-')[0] # training functions if d_aug: # https://github.com/mit-han-lab/data-efficient-gans train = EasyDict( run_func_name='training.training_loop_diffaug.training_loop' ) # Options for training loop (Diff Augment method) loss_args = EasyDict( func_name='training.loss_diffaug.ns_DiffAugment_r1', policy=diffaug_policy) # Options for loss (Diff Augment method) else: # original nvidia train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop (original from NVidia) G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg' ) # Options for generator loss. D_loss = EasyDict(func_name='training.loss.D_logistic_r1' ) # Options for discriminator loss. # network functions G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). G.impl = D.impl = ops # resolutions data_res = basename(tfr_file).split('-')[-1].split( 'x') # get resolution from dataset filename data_res = list(reversed([int(x) for x in data_res])) # convert to int list init_res, resolution, res_log2 = calc_init_res(data_res) if init_res != [4, 4]: print(' custom init resolution', init_res) G.init_res = D.init_res = list(init_res) train.setname = desc + config desc = '%s-%d-%s' % (desc, resolution, config) # training schedule sched.lod_training_kimg = lod_step_kimg sched.lod_transition_kimg = lod_step_kimg train.total_kimg = lod_step_kimg * res_log2 * 2 # a la ProGAN if finetune: train.total_kimg = 15000 # should start from ~10k kimg train.image_snapshot_ticks = 1 train.network_snapshot_ticks = 5 train.mirror_augment = mirror train.mirror_augment_v = mirror_v # learning rate if config == 'e': if finetune: # uptrain 1024 sched.G_lrate_base = 0.001 else: # train 1024 sched.G_lrate_base = 0.001 sched.G_lrate_dict = {0: 0.001, 1: 0.0007, 2: 0.0005, 3: 0.0003} sched.lrate_step = 1500 # period for stepping to next lrate, in kimg if config == 'f': # sched.G_lrate_base = 0.0003 sched.G_lrate_base = 0.001 sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway sched.minibatch_gpu_base = batch_size sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base sc.num_gpus = num_gpus if config == 'e': G.fmap_base = D.fmap_base = 8 << 10 if d_aug: loss_args.gamma = 100 if gamma is None else gamma else: D_loss.gamma = 100 if gamma is None else gamma elif config == 'f': G.fmap_base = D.fmap_base = 16 << 10 else: print(' Only configs E and F are implemented') exit() if cond: desc += '-cond' dataset_args.max_label_size = 'full' # conditioned on full label if freezeD: D.freezeD = True train.resume_with_new_nets = True if d_aug: desc += '-daug' sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, tf_config=tf_config) kwargs.update(resume_pkl=resume, resume_kimg=resume_kimg, resume_with_new_nets=True) if ema_kimg is not None: kwargs.update(G_ema_kimg=ema_kimg) if d_aug: kwargs.update(loss_args=loss_args) else: kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = train_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def run(g_loss, g_loss_kwargs, d_loss, d_loss_kwargs, dataset_train, dataset_eval, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics, resume_pkl, resume_kimg, resume_pkl_dir, max_images, lrate_base, img_ticks, net_ticks, skip_images): if g_loss_kwargs != '': g_loss_kwargs = json.loads(g_loss_kwargs) else: g_loss_kwargs = {} if d_loss_kwargs != '': d_loss_kwargs = json.loads(d_loss_kwargs) else: d_loss_kwargs = {} train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop. G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. G_loss = EasyDict( func_name='training.loss.' + g_loss, **g_loss_kwargs ) #G_logistic_ns_gsreg') # Options for generator loss. D_loss = EasyDict(func_name='training.loss.' + d_loss, **d_loss_kwargs) # Options for discriminator loss. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). train.total_kimg = total_kimg train.mirror_augment = mirror_augment train.image_snapshot_ticks = img_ticks train.network_snapshot_ticks = net_ticks G.scale_func = 'training.networks_stylegan2.apply_identity' D.scale_func = None sched.G_lrate_base = sched.D_lrate_base = lrate_base #0.002 # TODO: Changed this to 16 to match DiffAug sched.minibatch_size_base = 16 sched.minibatch_gpu_base = 4 D_loss.gamma = 10 metrics = [metric_defaults[x] for x in metrics] desc = 'stylegan2' sched.tick_kimg_base = 1 sched.tick_kimg_dict = { } #{8:28, 16:24, 32:20, 64:16, 128:12, 256:8, 512:6, 1024:4}): # Resolution-specific overrides. desc += '-' + dataset_train.split('/')[-1] # Get dataset paths t_path = dataset_train.split('/') e_path = dataset_eval.split('/') if len(t_path) > 1: dataset_train = t_path[-1] train.train_data_dir = os.path.join(data_dir, '/'.join(t_path[:-1])) if len(e_path) > 1: dataset_eval = e_path[-1] train.eval_data_dir = os.path.join(data_dir, '/'.join(e_path[:-1])) dataset_args = EasyDict(tfrecord_dir=dataset_train) # Limit number of training images during train (not eval) dataset_args['max_images'] = max_images if max_images: desc += '-%dimg' % max_images dataset_args['skip_images'] = skip_images dataset_args_eval = EasyDict(tfrecord_dir=dataset_eval) desc += '-' + dataset_eval assert num_gpus in [1, 2, 4, 8] sc.num_gpus = num_gpus assert config_id in _valid_configs desc += '-' + config_id if mirror_augment: desc += '-aug' # Infer pretrain checkpoint from target dataset if not resume_pkl: if any(ds in dataset_train.lower() for ds in ['obama', 'celeba', 'rem', 'portrait']): resume_pkl = 'ffhq-config-f.pkl' if any(ds in dataset_train.lower() for ds in ['gogh', 'temple', 'tower', 'medici', 'bridge']): resume_pkl = 'church-config-f.pkl' if any(ds in dataset_train.lower() for ds in ['bus']): resume_pkl = 'car-config-f.pkl' resume_pkl = os.path.join(resume_pkl_dir, resume_pkl) train.resume_pkl = resume_pkl train.resume_kimg = resume_kimg train.resume_with_new_nets = True # Recreate with new parameters # Adaptive parameters if 'ada' in config_id: G['train_scope'] = D[ 'train_scope'] = '.*adapt' # Freeze old parameters if 'ss' in config_id: G['adapt_func'] = D[ 'adapt_func'] = 'training.networks_stylegan2.apply_adaptive_scale_shift' if 'sv' or 'pc' in config_id: # [:9] == 'config-sv' or config_id[:9] == 'config-pc': G['map_svd'] = G['syn_svd'] = D['svd'] = True # Flatten over spatial dimension if 'flat' in config_id: G['spatial'] = D['spatial'] = True # Do PCA by centering before SVD if 'pc' in config_id: G['svd_center'] = D['svd_center'] = True G['svd_config'] = D['svd_config'] = 'S' if 'U' in config_id: G['svd_config'] += 'U' D['svd_config'] += 'U' if 'V' in config_id: G['svd_config'] += 'V' D['svd_config'] += 'V' # FreezeD D['freeze'] = 'fd' in config_id #freeze_d # DiffAug if 'da' in config_id: G_loss = EasyDict(func_name='training.loss.G_ns_diffaug') D_loss = EasyDict(func_name='training.loss.D_ns_diffaug_r1') if gamma is not None: D_loss.gamma = gamma sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update( G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss, ) kwargs.update(dataset_args=dataset_args, dataset_args_eval=dataset_args_eval, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def run(data, train_dir, config, d_aug, diffaug_policy, cond, ops, mirror, mirror_v, \ kimg, batch_size, lrate, resume, resume_kimg, num_gpus, ema_kimg, gamma, freezeD): # training functions if d_aug: # https://github.com/mit-han-lab/data-efficient-gans train = EasyDict( run_func_name='training.training_loop_diffaug.training_loop' ) # Options for training loop (Diff Augment method) loss_args = EasyDict( func_name='training.loss_diffaug.ns_DiffAugment_r1', policy=diffaug_policy) # Options for loss (Diff Augment method) else: # original nvidia train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop (original from NVidia) G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg' ) # Options for generator loss. D_loss = EasyDict(func_name='training.loss.D_logistic_r1' ) # Options for discriminator loss. # network functions G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). G.impl = D.impl = ops # dataset (tfrecords) - get or create tfr_files = file_list(os.path.dirname(data), 'tfr') tfr_files = [ f for f in tfr_files if basename(data) == basename(f).split('-')[0] ] if len(tfr_files) == 0 or os.stat(tfr_files[0]).st_size == 0: tfr_file, total_samples = create_from_image_folders( data) if cond is True else create_from_images(data) else: tfr_file = tfr_files[0] dataset_args = EasyDict(tfrecord=tfr_file) # resolutions with tf.Graph().as_default(), tflib.create_session().as_default(): # pylint: disable=not-context-manager dataset_obj = dataset.load_dataset( **dataset_args) # loading the data to see what comes out resolution = dataset_obj.resolution init_res = dataset_obj.init_res res_log2 = dataset_obj.res_log2 dataset_obj.close() dataset_obj = None if list(init_res) == [4, 4]: desc = '%s-%d' % (basename(data), resolution) else: print(' custom init resolution', init_res) desc = basename(tfr_file) G.init_res = D.init_res = list(init_res) train.savenames = [desc.replace(basename(data), 'snapshot'), desc] desc += '-%s' % config # training schedule train.total_kimg = kimg train.image_snapshot_ticks = 1 * num_gpus if kimg <= 1000 else 4 * num_gpus train.network_snapshot_ticks = 5 train.mirror_augment = mirror train.mirror_augment_v = mirror_v sched.tick_kimg_base = 2 if train.total_kimg < 2000 else 4 # learning rate if config == 'e': sched.G_lrate_base = 0.001 sched.G_lrate_dict = {0: 0.001, 1: 0.0007, 2: 0.0005, 3: 0.0003} sched.lrate_step = 1500 # period for stepping to next lrate, in kimg if config == 'f': sched.G_lrate_base = lrate # 0.001 for big datasets, 0.0003 for few-shot sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway # batch size (for 16gb memory GPU) sched.minibatch_gpu_base = 4096 // resolution if batch_size is None else batch_size print(' Batch size', sched.minibatch_gpu_base) sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base sc.num_gpus = num_gpus if config == 'e': G.fmap_base = D.fmap_base = 8 << 10 if d_aug: loss_args.gamma = 100 if gamma is None else gamma else: D_loss.gamma = 100 if gamma is None else gamma elif config == 'f': G.fmap_base = D.fmap_base = 16 << 10 else: print(' Only configs E and F are implemented') exit() if cond: desc += '-cond' dataset_args.max_label_size = 'full' # conditioned on full label if freezeD: D.freezeD = True train.resume_with_new_nets = True if d_aug: desc += '-daug' sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, tf_config=tf_config) kwargs.update(resume_pkl=resume, resume_kimg=resume_kimg, resume_with_new_nets=True) if ema_kimg is not None: kwargs.update(G_ema_kimg=ema_kimg) if d_aug: kwargs.update(loss_args=loss_args) else: kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = train_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics): train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop. G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg' ) # Options for generator loss. D_loss = EasyDict(func_name='training.loss.D_logistic_r1' ) # Options for discriminator loss. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='8k', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). train.data_dir = data_dir train.total_kimg = total_kimg train.mirror_augment = mirror_augment train.image_snapshot_ticks = train.network_snapshot_ticks = 10 sched.G_lrate_base = sched.D_lrate_base = 0.002 sched.minibatch_size_base = 32 sched.minibatch_gpu_base = 4 D_loss.gamma = 10 metrics = [metric_defaults[x] for x in metrics] desc = 'stylegan2' desc += '-' + dataset dataset_args = EasyDict(tfrecord_dir=dataset) assert num_gpus in [1, 2, 4, 8] sc.num_gpus = num_gpus desc += '-%dgpu' % num_gpus assert config_id in _valid_configs desc += '-' + config_id # Configs A-E: Shrink networks to match original StyleGAN. if config_id not in ['config-f', 'config-l']: G.fmap_base = D.fmap_base = 8 << 10 # Config L: Generator training only if config_id == 'config-l': # Use labels as latent vector input dataset_args.max_label_size = "full" # Deactivate methods specific for GAN training G.truncation_psi = None G.randomize_noise = False G.style_mixing_prob = None G.dlatent_avg_beta = None G.conditional_labels = False # Refinement training G_loss.func_name = 'training.loss.G_reconstruction' train.run_func_name = 'training.training_loop.training_loop_refinement' # G.freeze_layers = ["mapping", "noise"]#, "4x4", "8x8", "16x16", "32x32"] # Network for refinement train.resume_pkl = "nets/stylegan2-ffhq-config-f.pkl" # TODO init net train.resume_with_new_nets = True # Maintenance tasks sched.tick_kimg_base = 1 # 1 tick = 5000 images (metric update) sched.tick_kimg_dict = {} train.image_snapshot_ticks = 5 # Save every 5000 images train.network_snapshot_ticks = 10 # Save every 10000 images # Training parameters sched.G_lrate_base = 1e-4 train.G_smoothing_kimg = 0.0 sched.minibatch_size_base = sched.minibatch_gpu_base * num_gpus # 4 per GPU # Config E: Set gamma to 100 and override G & D architecture. if config_id.startswith('config-e'): D_loss.gamma = 100 if 'Gorig' in config_id: G.architecture = 'orig' if 'Gskip' in config_id: G.architecture = 'skip' # (default) if 'Gresnet' in config_id: G.architecture = 'resnet' if 'Dorig' in config_id: D.architecture = 'orig' if 'Dskip' in config_id: D.architecture = 'skip' if 'Dresnet' in config_id: D.architecture = 'resnet' # (default) # Configs A-D: Enable progressive growing and switch to networks that support it. if config_id in ['config-a', 'config-b', 'config-c', 'config-d']: sched.lod_initial_resolution = 8 sched.G_lrate_base = sched.D_lrate_base = 0.001 sched.G_lrate_dict = sched.D_lrate_dict = { 128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003 } sched.minibatch_size_base = 32 # (default) sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32} sched.minibatch_gpu_base = 4 # (default) sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4} G.synthesis_func = 'G_synthesis_stylegan_revised' D.func_name = 'training.networks_stylegan2.D_stylegan' # Configs A-C: Disable path length regularization. if config_id in ['config-a', 'config-b', 'config-c']: G_loss = EasyDict(func_name='training.loss.G_logistic_ns') # Configs A-B: Disable lazy regularization. if config_id in ['config-a', 'config-b']: train.lazy_regularization = False # Config A: Switch to original StyleGAN networks. if config_id == 'config-a': G = EasyDict(func_name='training.networks_stylegan.G_style') D = EasyDict(func_name='training.networks_stylegan.D_basic') if gamma is not None: D_loss.gamma = gamma sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
r1_gamma=10.0) # Options for discriminator loss. dataset = EasyDict( tfrecord_dir='streetviewtf') # Options for load_dataset(). sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='4k', layout='random') # Options for setup_snapshot_image_grid(). #metrics = [metric_base.fid50k] # Options for MetricGroup. metrics = [] # Options for MetricGroup. submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). train.resume_snapshot = '/root/repos/idinvert/results/00016-sgan-1gpu/network-snapshot-015040.pkl' # Network pickle to resume training from, None = train from scratch. train.resume_kimg = 15000 # Assumed training progress at the beginning. Affects reporting and training schedule. train.resume_time = 0.0 # Assumed wallclock time at the beginning. Affects reporting. train.mirror_augment = True train.image_snapshot_ticks = 1 train.network_snapshot_ticks = 2 # Dataset. #desc += '-ffhq'; dataset = EasyDict(tfrecord_dir='ffhq'); train.mirror_augment = True ##desc += '-ffhq512'; dataset = EasyDict(tfrecord_dir='ffhq', resolution=512); train.mirror_augment = True #desc += '-ffhq256'; dataset = EasyDict(tfrecord_dir='ffhq', resolution=256); train.mirror_augment = True #desc += '-celebahq'; dataset = EasyDict(tfrecord_dir='celebahq'); train.mirror_augment = True #desc += '-bedroom'; dataset = EasyDict(tfrecord_dir='lsun-bedroom-full'); train.mirror_augment = False #desc += '-car'; dataset = EasyDict(tfrecord_dir='lsun-car-512x384'); train.mirror_augment = False #desc += '-cat'; dataset = EasyDict(tfrecord_dir='lsun-cat-full'); train.mirror_augment = False # Number of GPUs. desc += '-1gpu' submit_config.num_gpus = 1 sched.minibatch_base = 16 #; sched.minibatch_dict = {4: 128, 8: 128, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8, 512: 4}
32: 64, 64: 32, 128: 16, 256: 8, 512: 4 } #desc += '-2gpu'; submit_config.num_gpus = 2; sched.minibatch_base = 8; sched.minibatch_dict = {4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8} #desc += '-4gpu'; submit_config.num_gpus = 4; sched.minibatch_base = 16; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16} #desc += '-8gpu'; submit_config.num_gpus = 8; sched.minibatch_base = 32; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32} # Default options. train.total_kimg = args.total_kimg sched.lod_initial_resolution = args.init_res sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003} sched.D_lrate_dict = EasyDict(sched.G_lrate_dict) train.image_snapshot_ticks = args.nsteps_image_snapshot train.network_snapshot_ticks = args.nsteps_network_snapshot train.resume_run_id = args.resume_run_id train.resume_snapshot = args.resume_snapshot # WGAN-GP loss for CelebA-HQ. #desc += '-wgangp'; G_loss = EasyDict(func_name='training.loss.G_wgan'); D_loss = EasyDict(func_name='training.loss.D_wgan_gp'); sched.G_lrate_dict = {k: min(v, 0.002) for k, v in sched.G_lrate_dict.items()}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict) # Table 1. #desc += '-tuned-baseline'; G.use_styles = False; G.use_pixel_norm = True; G.use_instance_norm = False; G.mapping_layers = 0; G.truncation_psi = None; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False #desc += '-add-mapping-and-styles'; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False #desc += '-remove-traditional-input'; G.style_mixing_prob = 0.0; G.use_noise = False #desc += '-add-noise-inputs'; G.style_mixing_prob = 0.0 #desc += '-mixing-regularization' # default # Table 2.