def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, mirror_augment_v, metrics, min_h, min_w, res_log2, lr, use_attention, resume_with_new_nets, glr, dlr, use_raw, resume_pkl, minibatch_gpu_base, network_snapshot_ticks): 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.mirror_augment_v = mirror_augment_v train.resume_with_new_nets = resume_with_new_nets train.image_snapshot_ticks = 1 train.network_snapshot_ticks = network_snapshot_ticks sched.G_lrate_base = sched.D_lrate_base = lr train.resume_pkl = resume_pkl if glr: sched.G_lrate_base = glr if dlr: sched.D_lrate_base = dlr sched.minibatch_size_base = 32 sched.minibatch_gpu_base = minibatch_gpu_base D_loss.gamma = 10 metrics = [metric_defaults[x] for x in metrics] desc = 'stylegan2' desc += '-' + dataset dataset_args = EasyDict(tfrecord_dir=dataset) dataset_args.use_raw = use_raw G.min_h = D.min_h = dataset_args.min_h = min_h G.min_w = D.min_w = dataset_args.min_w = min_w G.res_log2 = D.res_log2 = dataset_args.res_log2 = res_log2 if use_attention: desc += '-attention' G.use_attention = True D.use_attention = True 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, data_dir, result_dir, num_gpus, total_kimg, mirror_augment, metrics, resume_pkl, model_type='vc_gan2', latent_type='uniform', batch_size=32, batch_per_gpu=16, random_seed=1000, G_fmap_base=8, module_G_list=None, G_nf_scale=4, E_fmap_base=8, module_E_list=None, E_nf_scale=4, D_fmap_base=9, module_D_list=None, D_nf_scale=4, I_fmap_base=9, module_I_list=None, I_nf_scale=4, fmap_decay=0.15, fmap_min=16, fmap_max=512, n_samples_per=10, topk_dims_to_show=20, hy_beta=1, hy_gamma=0, hy_dcp=40, hy_ncut=1, hy_rec=20, hy_hes=20, hy_lin=20, hy_mat=80, hy_gmat=0, hy_oth=80, hy_det=0, hessian_type='no_act_points', n_act_points=10, lie_alg_init_type='oth', lie_alg_init_scale=0.1, G_lrate_base=0.002, D_lrate_base=None, lambda_d_factor=10., lambda_od=1., group_loss_type='_rec_mat_', group_feats_size=400, temp=0.67, n_discrete=0, drange_net=[-1, 1], recons_type='bernoulli_loss', use_group_decomp=False, snapshot_ticks=10): train = EasyDict( run_func_name='training.training_loop_vae.training_loop_vae' ) # Options for training loop. if not (module_G_list is None): module_G_list = _str_to_list(module_G_list) key_G_ls, size_G_ls, count_dlatent_G_size = split_module_names( module_G_list) if not (module_E_list is None): module_E_list = _str_to_list(module_E_list) key_E_ls, size_E_ls, count_dlatent_E_size = split_module_names( module_E_list) if not (module_D_list is None): module_D_list = _str_to_list(module_D_list) key_D_ls, size_D_ls, count_dlatent_D_size = split_module_names( module_D_list) if not (module_I_list is None): module_I_list = _str_to_list(module_I_list) key_I_ls, size_I_ls, count_dlatent_I_size = split_module_names( module_I_list) D = D_opt = D_loss = None E = EasyDict(func_name='training.vae_networks.E_main_modular', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_E_size, group_feats_size=group_feats_size, module_E_list=module_E_list, nf_scale=E_nf_scale, n_discrete=n_discrete, fmap_base=2 << E_fmap_base) # Options for encoder network. G = EasyDict(func_name='training.vae_networks.G_main_modular', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_G_size, group_feats_size=group_feats_size, module_G_list=module_G_list, nf_scale=G_nf_scale, n_discrete=n_discrete, recons_type=recons_type, n_act_points=n_act_points, lie_alg_init_type=lie_alg_init_type, lie_alg_init_scale=lie_alg_init_scale, fmap_base=2 << G_fmap_base) # Options for generator network. I = EasyDict(func_name='training.vae_I_networks.I_main_modular', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_I_size, module_I_list=module_I_list, nf_scale=I_nf_scale, fmap_base=2 << I_fmap_base) # Options for I network. G_opt = EasyDict(beta1=0.9, beta2=0.999, epsilon=1e-8) # Options for generator optimizer. if model_type == 'factor_vae' or model_type == 'factor_sindis_vae': # Factor-VAE D = EasyDict( func_name='training.vae_networks.D_factor_vae_modular', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_D_size, module_D_list=module_D_list, nf_scale=D_nf_scale, fmap_base=2 << D_fmap_base) # Options for generator network. D_opt = EasyDict(beta1=0.5, beta2=0.9, epsilon=1e-8) # Options for discriminator optimizer. desc = model_type + '_modular' if model_type == 'beta_vae': # Beta-VAE G_loss = EasyDict( func_name='training.loss_vae.beta_vae', latent_type=latent_type, hy_beta=hy_beta, recons_type=recons_type) # Options for generator loss. elif model_type == 'betatc_vae': # BetaTC-VAE G_loss = EasyDict( func_name='training.loss_vae.betatc_vae', latent_type=latent_type, hy_beta=hy_beta, recons_type=recons_type) # Options for generator loss. elif model_type == 'lie_vae': # LieVAE G_loss = EasyDict( func_name='training.loss_vae_lie.lie_vae', latent_type=latent_type, hy_rec=hy_rec, hy_dcp=hy_dcp, hy_hes=hy_hes, hy_lin=hy_lin, hy_ncut=hy_ncut, recons_type=recons_type) # Options for generator loss. elif model_type == 'lie_vae_with_split': # LieVAE with split loss G_loss = EasyDict( func_name='training.loss_vae_lie.lie_vae_with_split', latent_type=latent_type, hy_rec=hy_rec, hy_dcp=hy_dcp, hy_hes=hy_hes, hy_lin=hy_lin, hy_ncut=hy_ncut, recons_type=recons_type) # Options for generator loss. elif model_type == 'group_vae_v2': # GroupVAE-v2 G_loss = EasyDict( func_name='training.loss_vae_group_v2.group_act_vae', latent_type=latent_type, hy_beta=hy_beta, hy_rec=hy_rec, hy_gmat=hy_gmat, hy_dcp=hy_dcp, hy_hes=hy_hes, hy_lin=hy_lin, hy_ncut=hy_ncut, hessian_type=hessian_type, recons_type=recons_type) # Options for generator loss. elif model_type == 'group_vae_spl_v2': # GroupVAE-v2 G_loss = EasyDict( func_name='training.loss_vae_group_v2.group_act_spl_vae', latent_type=latent_type, hy_beta=hy_beta, hy_rec=hy_rec, hy_gmat=hy_gmat, hy_dcp=hy_dcp, hy_hes=hy_hes, hy_lin=hy_lin, hy_ncut=hy_ncut, hessian_type=hessian_type, recons_type=recons_type) # Options for generator loss. elif model_type == 'group_vae': # Group-VAE G_loss = EasyDict( func_name='training.loss_vae.group_vae', latent_type=latent_type, hy_beta=hy_beta, hy_dcp=hy_dcp, hy_ncut=hy_ncut, hy_rec=hy_rec, hy_mat=hy_mat, hy_gmat=hy_gmat, hy_oth=hy_oth, hy_det=hy_det, use_group_decomp=use_group_decomp, group_loss_type=group_loss_type, recons_type=recons_type) # Options for generator loss. elif model_type == 'group_vae_wc': # Group-VAE-with_Cat G_loss = EasyDict( func_name='training.loss_vae.group_vae_wc', latent_type=latent_type, hy_beta=hy_beta, hy_gamma=hy_gamma, temp=temp, use_group_decomp=use_group_decomp, group_loss_type=group_loss_type, recons_type=recons_type) # Options for generator loss. elif model_type == 'dip_vae_i' or model_type == 'dip_vae_ii': # DIP-VAE G_loss = EasyDict( func_name='training.loss_vae.dip_vae', lambda_d_factor=lambda_d_factor, lambda_od=lambda_od, latent_type=latent_type, dip_type=model_type, recons_type=recons_type) # Options for generator loss. elif model_type == 'factor_vae': # Factor-VAE G_loss = EasyDict( func_name='training.loss_vae.factor_vae_G', latent_type=latent_type, hy_gamma=hy_gamma, recons_type=recons_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vae.factor_vae_D', latent_type=latent_type) # Options for discriminator loss. elif model_type == 'factor_sindis_vae': # Factor-VAE G_loss = EasyDict( func_name='training.loss_vae.factor_vae_sindis_G', latent_type=latent_type, hy_gamma=hy_gamma, recons_type=recons_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vae.factor_vae_sindis_D', latent_type=latent_type) # 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': random_seed, 'allow_soft_placement': True } # 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 = snapshot_ticks sched.G_lrate_base = G_lrate_base sched.D_lrate_base = D_lrate_base sched.minibatch_size_base = batch_size sched.minibatch_gpu_base = batch_per_gpu metrics = [metric_defaults[x] for x in metrics] desc += '-' + dataset dataset_args = EasyDict(tfrecord_dir=dataset, max_label_size='full') assert num_gpus in [1, 2, 4, 8] sc.num_gpus = num_gpus desc += '-%dgpu' % num_gpus sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, E_args=E, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss, traversal_grid=True) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, n_continuous=count_dlatent_G_size, n_discrete=n_discrete, drange_net=drange_net, metric_arg_list=metrics, tf_config=tf_config, resume_pkl=resume_pkl, n_samples_per=n_samples_per, topk_dims_to_show=topk_dims_to_show) 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(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 = float( os.environ['G_LR']) if 'G_LR' in os.environ else 0.002 sched.D_lrate_base = float( os.environ['D_LR']) if 'D_LR' in os.environ else 0.002 sched.G_lrate_base *= float( os.environ['G_LR_MULT']) if 'G_LR_MULT' in os.environ else 1.0 sched.D_lrate_base *= float( os.environ['D_LR_MULT']) if 'D_LR_MULT' in os.environ else 1.0 G_opt.beta2 = float( os.environ['G_BETA2']) if 'G_BETA2' in os.environ else 0.99 D_opt.beta2 = float( os.environ['D_BETA2']) if 'D_BETA2' in os.environ else 0.99 print('G_lrate: %f' % sched.G_lrate_base) print('D_lrate: %f' % sched.D_lrate_base) print('G_beta2: %f' % G_opt.beta2) print('D_beta2: %f' % D_opt.beta2) sched.minibatch_size_base = int( os.environ['BATCH_SIZE']) if 'BATCH_SIZE' in os.environ else num_gpus sched.minibatch_gpu_base = int( os.environ['BATCH_PER']) if 'BATCH_PER' in os.environ else 1 D_loss.gamma = 10 metrics = [metric_defaults[x] for x in metrics] desc = 'stylegan2' desc += '-' + dataset resolution = int( os.environ['RESOLUTION']) if 'RESOLUTION' in os.environ else 64 dataset_args = EasyDict(tfrecord_dir=dataset, resolution=resolution) assert num_gpus in [ 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192 ] 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 if 'FMAP_BASE' in os.environ: G.fmap_base = D.fmap_base = int(os.environ['FMAP_BASE']) << 10 else: G.fmap_base = D.fmap_base = 16 << 10 # default print('G_fmap_base: %d' % G.fmap_base) print('D_fmap_base: %d' % D.fmap_base) # 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, data_dir, result_dir, num_gpus, total_kimg, mirror_augment, metrics, resume_pkl, model_type='vc_gan2', latent_type='uniform', batch_size=32, batch_per_gpu=16, random_seed=1000, G_fmap_base=8, module_G_list=None, G_nf_scale=4, E_fmap_base=8, module_E_list=None, E_nf_scale=4, D_fmap_base=9, module_D_list=None, D_nf_scale=4, fmap_decay=0.15, fmap_min=16, fmap_max=512, n_samples_per=10, arch='resnet', topk_dims_to_show=20, hy_beta=1, hy_gamma=0, hy_1p=0, lie_alg_init_type='oth', lie_alg_init_scale=0.1, G_lrate_base=0.002, D_lrate_base=None, group_feats_size=400, temp=0.67, n_discrete=0, epsilon=1, drange_net=[-1, 1], recons_type='bernoulli_loss', R_view_scale=1, group_feat_type='concat', use_sphere_points=False, use_learnable_sphere_points=False, n_sphere_points=100, mapping_after_exp=False, snapshot_ticks=10): train = EasyDict( run_func_name='training.training_loop_gan.training_loop_gan' ) # Options for training loop. if not (module_G_list is None): module_G_list = _str_to_list(module_G_list) key_G_ls, size_G_ls, count_dlatent_G_size = split_module_names( module_G_list) if not (module_E_list is None): module_E_list = _str_to_list(module_E_list) key_E_ls, size_E_ls, count_dlatent_E_size = split_module_names( module_E_list) if not (module_D_list is None): module_D_list = _str_to_list(module_D_list) key_D_ls, size_D_ls, count_dlatent_D_size = split_module_names( module_D_list) E = EasyDict(func_name='training.gan_networks.E_main_modular', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_E_size, group_feats_size=group_feats_size, module_E_list=module_E_list, nf_scale=E_nf_scale, n_discrete=n_discrete, fmap_base=2 << E_fmap_base) # Options for encoder network. D = EasyDict( func_name='training.gan_networks.D_main_modular', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_D_size, group_feats_size=group_feats_size, module_D_list=module_D_list, nf_scale=D_nf_scale, n_discrete=n_discrete, fmap_base=2 << D_fmap_base) # Options for discriminator network. G = EasyDict(func_name='training.gan_networks.G_main_modular', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_G_size, group_feats_size=group_feats_size, module_G_list=module_G_list, nf_scale=G_nf_scale, n_discrete=n_discrete, recons_type=recons_type, lie_alg_init_type=lie_alg_init_type, lie_alg_init_scale=lie_alg_init_scale, R_view_scale=R_view_scale, group_feat_type=group_feat_type, mapping_after_exp=mapping_after_exp, use_sphere_points=use_sphere_points, use_learnable_sphere_points=use_learnable_sphere_points, n_sphere_points=n_sphere_points, fmap_base=2 << G_fmap_base) # Options for generator network. G_opt = EasyDict(beta1=0.9, beta2=0.999, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.9, beta2=0.999, epsilon=1e-8) # Options for discriminator optimizer. desc = model_type + '_modular' if model_type == 'so_gan': G_loss = EasyDict( func_name='training.loss_gan_so.so_gan', hy_1p=hy_1p, hy_beta=hy_beta, latent_type=latent_type, recons_type=recons_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_gan.gan_D', latent_type=latent_type) # Options for discriminator loss. else: raise ValueError('Unknown model_type:', model_type) 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': random_seed, 'allow_soft_placement': True } # 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 = snapshot_ticks sched.G_lrate_base = G_lrate_base sched.D_lrate_base = D_lrate_base sched.minibatch_size_base = batch_size sched.minibatch_gpu_base = batch_per_gpu metrics = [metric_defaults[x] for x in metrics] desc += '-' + dataset dataset_args = EasyDict(tfrecord_dir=dataset, max_label_size='full') assert num_gpus in [1, 2, 4, 8] sc.num_gpus = num_gpus desc += '-%dgpu' % num_gpus sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, E_args=E, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss, traversal_grid=True) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, n_continuous=count_dlatent_G_size, n_discrete=n_discrete, drange_net=drange_net, metric_arg_list=metrics, tf_config=tf_config, resume_pkl=resume_pkl, n_samples_per=n_samples_per, topk_dims_to_show=topk_dims_to_show) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)