def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics, resume_G_pkl=None, n_batch=2, n_batch_per_gpu=1, D_global_size=0, C_global_size=10, model_type='hd_dis_model', latent_type='uniform', resume_pkl=None, n_samples_per=4, D_lambda=0, C_lambda=1, epsilon_in_loss=3, random_eps=True, M_lrmul=0.1, resolution_manual=1024, pretrained_type='with_stylegan2', traj_lambda=None, level_I_kimg=1000, use_level_training=False, resume_kimg=0, use_std_in_m=False, prior_latent_size=512, stylegan2_dlatent_size=512, stylegan2_mapping_fmaps=512, M_mapping_fmaps=512, hyperplane_lambda=1, hyperdir_lambda=1): train = EasyDict( run_func_name='training.training_loop_hdwG.training_loop_hdwG') G = EasyDict(func_name='training.hd_networks_stylegan2.G_main', latent_size=prior_latent_size, dlatent_size=stylegan2_dlatent_size, mapping_fmaps=stylegan2_mapping_fmaps, mapping_lrmul=M_lrmul, style_mixing_prob=None, dlatent_avg_beta=None, truncation_psi=None, normalize_latents=False, structure='fixed') D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2') if model_type == 'hd_hyperplane': M = EasyDict(func_name='training.hd_networks.net_M_hyperplane', C_global_size=C_global_size, D_global_size=D_global_size, latent_size=prior_latent_size, mapping_lrmul=M_lrmul, use_std_in_m=use_std_in_m) I = EasyDict(func_name='training.hd_networks.net_I', C_global_size=C_global_size, D_global_size=D_global_size) elif model_type == 'vc_gan_preprior': M = EasyDict(func_name='training.hd_networks.net_M_vc', C_global_size=C_global_size, D_global_size=D_global_size, latent_size=prior_latent_size, mapping_lrmul=M_lrmul, use_std_in_m=use_std_in_m) I = EasyDict(func_name='training.hd_networks.net_I', C_global_size=C_global_size, D_global_size=D_global_size) elif model_type == 'vc_gan': M = EasyDict(func_name='training.hd_networks.net_M_empty', C_global_size=C_global_size, D_global_size=D_global_size, latent_size=prior_latent_size, mapping_lrmul=M_lrmul, use_std_in_m=use_std_in_m) I = EasyDict(func_name='training.hd_networks.net_I', C_global_size=C_global_size, D_global_size=D_global_size) G.mapping_func = 'G_mapping_hd_dis_to_dlatent' else: M = EasyDict(func_name='training.hd_networks.net_M', C_global_size=C_global_size, D_global_size=D_global_size, latent_size=prior_latent_size, mapping_fmaps=M_mapping_fmaps, mapping_lrmul=M_lrmul, use_std_in_m=use_std_in_m) I = EasyDict(func_name='training.hd_networks.net_I', C_global_size=C_global_size, D_global_size=D_global_size) if model_type == 'hd_dis_model_with_cls': I_info = EasyDict(func_name='training.hd_networks.net_I_info', C_global_size=C_global_size, D_global_size=D_global_size) else: I_info = EasyDict() I_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) if model_type == 'vc_gan': I_loss = EasyDict(func_name='training.loss_hdwG.IandG_vc_loss', latent_type=latent_type, D_global_size=D_global_size, C_global_size=C_global_size, D_lambda=D_lambda, C_lambda=C_lambda, epsilon=epsilon_in_loss, random_eps=random_eps, traj_lambda=traj_lambda, resolution_manual=resolution_manual, use_std_in_m=use_std_in_m, model_type=model_type, hyperplane_lambda=hyperplane_lambda, prior_latent_size=prior_latent_size, hyperdir_lambda=hyperdir_lambda) else: I_loss = EasyDict( func_name='training.loss_hdwG.IandMandG_hyperplane_loss', latent_type=latent_type, D_global_size=D_global_size, C_global_size=C_global_size, D_lambda=D_lambda, C_lambda=C_lambda, epsilon=epsilon_in_loss, random_eps=random_eps, traj_lambda=traj_lambda, resolution_manual=resolution_manual, use_std_in_m=use_std_in_m, model_type=model_type, hyperplane_lambda=hyperplane_lambda, prior_latent_size=prior_latent_size, hyperdir_lambda=hyperdir_lambda) D_loss = EasyDict(func_name='training.loss.D_logistic_r1') sched = EasyDict() grid = EasyDict(size='1080p', layout='random') sc = dnnlib.SubmitConfig() tf_config = {'rnd.np_random_seed': 1000} 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 = n_batch sched.minibatch_gpu_base = n_batch_per_gpu D_loss.gamma = 10 metrics = [metric_defaults[x] for x in metrics] desc = 'hdwG_disentanglement' 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 if resolution_manual <= 256: I.fmap_base = 2 << 8 G.fmap_base = 2 << 10 D.fmap_base = 2 << 8 else: I.fmap_base = 8 << 10 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.002 # sched.G_lrate_dict = sched.D_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003} sched.minibatch_size_base = n_batch # (default) # sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32} sched.minibatch_gpu_base = n_batch_per_gpu # (default) # sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4} # G.synthesis_func = 'hd_networks_stylegan2.G_synthesis_stylegan_revised' G.synthesis_func = 'G_synthesis_stylegan_revised_hd' # D.func_name = 'training.networks_stylegan2.D_stylegan' # 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, I_args=I, M_args=M, I_opt_args=I_opt, D_opt_args=D_opt, I_loss_args=I_loss, D_loss_args=D_loss, resume_pkl=resume_pkl, resume_G_pkl=resume_G_pkl) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, use_hd_with_cls=(model_type == 'hd_dis_model_with_cls'), use_hyperplane=(model_type == 'hd_hyperplane'), metric_arg_list=metrics, tf_config=tf_config, n_discrete=D_global_size, n_continuous=C_global_size, n_samples_per=n_samples_per, resolution_manual=resolution_manual, pretrained_type=pretrained_type, level_I_kimg=level_I_kimg, use_level_training=use_level_training, resume_kimg=resume_kimg, use_std_in_m=use_std_in_m, prior_latent_size=prior_latent_size, latent_type=latent_type) 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, gamma, mirror_augment, metrics, resume_pkl, I_fmap_base=8, G_fmap_base=8, D_fmap_base=9, fmap_decay=0.15, C_lambda=1, n_samples_per=10, module_list=None, model_type='gan_like', epsilon_loss=3, random_eps=False, latent_type='uniform', batch_size=32, batch_per_gpu=16, return_atts=False, random_seed=1000, module_I_list=None, module_D_list=None, fmap_min=16, fmap_max=512, G_nf_scale=4, norm_ord=2, topk_dims_to_show=20, learning_rate=0.002, avg_mv_for_I=False, use_cascade=False, cascade_alt_freq_k=1, post_trans_wh=16, post_trans_cnn_dim=128, dff=512, trans_rate=0.1, construct_feat_by_concat=False, ncut_maxval=5, post_trans_mat=16, group_recons_lambda=0, trans_dim=512, network_snapshot_ticks=10): train = EasyDict( run_func_name='training.training_loop_tsfm.training_loop_tsfm' ) # Options for training loop. if not (module_list is None): module_list = _str_to_list(module_list) key_ls, size_ls, count_dlatent_size = split_module_names(module_list) if model_type == 'info_gan_like': # Independent branch version InfoGAN G = EasyDict( func_name='training.tsfm_G_nets.G_main_tsfm', synthesis_func='G_synthesis_modular_tsfm', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale, trans_dim=trans_dim, post_trans_wh=post_trans_wh, post_trans_cnn_dim=post_trans_cnn_dim, dff=dff, trans_rate=trans_rate, construct_feat_by_concat=construct_feat_by_concat, ncut_maxval=ncut_maxval, post_trans_mat=post_trans_mat, ) # Options for generator network. I = EasyDict(func_name='training.tsfm_I_nets.head_infogan2', dlatent_size=count_dlatent_size, fmap_min=fmap_min, fmap_max=fmap_max) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. desc = 'info_gan_like' elif model_type == 'ps_sc_like': # COMA-FAIN G = EasyDict( func_name='training.tsfm_G_nets.G_main_tsfm', synthesis_func='G_synthesis_modular_tsfm', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale, trans_dim=trans_dim, post_trans_wh=post_trans_wh, post_trans_cnn_dim=post_trans_cnn_dim, dff=dff, trans_rate=trans_rate, construct_feat_by_concat=construct_feat_by_concat, ncut_maxval=ncut_maxval, post_trans_mat=post_trans_mat, ) # Options for generator network. I = EasyDict(func_name='training.tsfm_I_nets.head_ps_sc', dlatent_size=count_dlatent_size, fmap_min=fmap_min, fmap_max=fmap_max) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. desc = 'ps_sc_like' elif model_type == 'gan_like': # Just modular GAN. G = EasyDict( func_name='training.tsfm_G_nets.G_main_tsfm', synthesis_func='G_synthesis_modular_tsfm', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale, trans_dim=trans_dim, post_trans_wh=post_trans_wh, post_trans_cnn_dim=post_trans_cnn_dim, dff=dff, trans_rate=trans_rate, construct_feat_by_concat=construct_feat_by_concat, ncut_maxval=ncut_maxval, post_trans_mat=post_trans_mat, ) # Options for generator network. I = EasyDict() D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. desc = 'gan_like' else: raise ValueError('Not supported model tyle: ' + model_type) 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. if model_type == 'info_gan_like': # InfoGAN G_loss = EasyDict( func_name='training.loss_tsfm.G_logistic_ns_info_gan', C_lambda=C_lambda, latent_type=latent_type, norm_ord=norm_ord, group_recons_lambda=group_recons_lambda ) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_tsfm.D_logistic_r1_shared', latent_type=latent_type) # Options for discriminator loss. elif model_type == 'ps_sc_like': # PS-SC G_loss = EasyDict( func_name='training.loss_tsfm.G_logistic_ns_ps_sc', C_lambda=C_lambda, group_recons_lambda=group_recons_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, use_cascade=use_cascade) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_tsfm.D_logistic_r1_shared', latent_type=latent_type) # Options for discriminator loss. elif model_type == 'gan_like': # Just GAN G_loss = EasyDict(func_name='training.loss_tsfm.G_logistic_ns', latent_type=latent_type, group_recons_lambda=group_recons_lambda ) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_tsfm.D_logistic_r1_shared', 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 } # 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 = learning_rate 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 I.fmap_base = 2 << I_fmap_base G.fmap_base = 2 << G_fmap_base D.fmap_base = 2 << D_fmap_base 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, I_args=I, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss, use_info_gan=( model_type == 'info_gan_like'), # Independent branch version use_ps_head=(model_type == 'ps_sc_like'), avg_mv_for_I=avg_mv_for_I, traversal_grid=True, return_atts=return_atts) n_continuous = 0 if not (module_list is None): for i, key in enumerate(key_ls): m_name = key.split('-')[0] if (m_name in LATENT_MODULES) and (not m_name == 'D_global'): n_continuous += size_ls[i] kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config, resume_pkl=resume_pkl, n_continuous=n_continuous, n_samples_per=n_samples_per, topk_dims_to_show=topk_dims_to_show, cascade_alt_freq_k=cascade_alt_freq_k, network_snapshot_ticks=network_snapshot_ticks) 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(result_dir, num_gpus, total_kimg, mirror_augment, metrics, resume_pkl, G_pkl, I_fmap_base=8, fmap_decay=0.15, n_samples_per=10, module_list=None, latent_type='uniform', batch_size=32, batch_per_gpu=16, random_seed=1000, fmap_min=16, fmap_max=512, dlatent_size=10, I_nf_scale=4, arch='resnet'): print('module_list:', module_list) train = EasyDict( run_func_name='training.training_loop_infernet.training_loop_infernet' ) # Options for training loop. module_list = _str_to_list(module_list) I = EasyDict(func_name='training.vc_networks2.infer_modular', dlatent_size=dlatent_size, fmap_min=fmap_min, fmap_max=fmap_max, module_list=module_list, I_nf_scale=I_nf_scale) desc = 'inference_net' I_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. loss = EasyDict(func_name='training.loss_inference.I_loss', latent_type=latent_type, dlatent_size=dlatent_size) # Options for generator loss. sched = EasyDict() # Options for TrainingSchedule. sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). # tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). tf_config = { 'rnd.np_random_seed': random_seed } # Options for tflib.init_tf(). train.total_kimg = total_kimg sched.lrate = 0.002 sched.tick_kimg = 1 sched.minibatch_size = batch_size sched.minibatch_gpu = batch_per_gpu metrics = [metric_defaults[x] for x in metrics] assert num_gpus in [1, 2, 4, 8] sc.num_gpus = num_gpus desc += '-%dgpu' % num_gpus # Configs A-E: Shrink networks to match original StyleGAN. I.fmap_base = 2 << I_fmap_base sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(I_args=I, I_opt_args=I_opt, loss_args=loss) kwargs.update(sched_args=sched, metric_arg_list=metrics, tf_config=tf_config, resume_pkl=resume_pkl, G_pkl=G_pkl, n_samples_per=n_samples_per) 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 generate_images(network_pkl, seeds, truncation_psi, data_dir=None, dataset_name=None, model=None): G_args = EasyDict(func_name='training.' + model + '.G_main') dataset_args = EasyDict(tfrecord_dir=dataset_name) G_args.fmap_base = 8 << 10 tflib.init_tf() training_set = dataset.load_dataset(data_dir=dnnlib.convert_path(data_dir), verbose=True, **dataset_args) print('Constructing networks...') Gs = tflib.Network('G', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **G_args) print('Loading networks from "%s"...' % network_pkl) _, _, _Gs = pretrained_networks.load_networks(network_pkl) Gs.copy_vars_from(_Gs) noise_vars = [ var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise') ] Gs_kwargs = dnnlib.EasyDict() # Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) Gs_kwargs.randomize_noise = False if truncation_psi is not None: Gs_kwargs.truncation_psi = truncation_psi for seed_idx, seed in enumerate(seeds): print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds))) rnd = np.random.RandomState(seed) z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component] tflib.set_vars( {var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width] images, x_v, n_v, m_v = Gs.run( z, None, **Gs_kwargs) # [minibatch, height, width, channel] print(images.shape, n_v.shape, x_v.shape, m_v.shape) misc.convert_to_pil_image(images[0], drange=[-1, 1]).save( dnnlib.make_run_dir_path('seed%04d.png' % seed)) misc.save_image_grid(adjust_range(n_v), dnnlib.make_run_dir_path('seed%04d-nv.png' % seed), drange=[-1, 1]) print(np.linalg.norm(x_v - m_v)) misc.save_image_grid(adjust_range(x_v).transpose([1, 0, 2, 3]), dnnlib.make_run_dir_path('seed%04d-xv.png' % seed), drange=[-1, 1]) misc.save_image_grid(adjust_range(m_v).transpose([1, 0, 2, 3]), dnnlib.make_run_dir_path('seed%04d-mv.png' % seed), drange=[-1, 1]) misc.save_image_grid(adjust_range(clip(x_v, 'cat')), dnnlib.make_run_dir_path('seed%04d-xvs.png' % seed), drange=[-1, 1]) misc.save_image_grid(adjust_range(clip(m_v, 'ss')), dnnlib.make_run_dir_path('seed%04d-mvs.png' % seed), drange=[-1, 1]) misc.save_image_grid(adjust_range(clip(m_v, 'ffhq')), dnnlib.make_run_dir_path('seed%04d-fmvs.png' % seed), drange=[-1, 1])
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics, resume_pkl, I_fmap_base=8, G_fmap_base=8, D_fmap_base=9, fmap_decay=0.15, D_lambda=0, C_lambda=1, n_samples_per=10, module_list=None, model_type='vc_gan2', epsilon_loss=3, random_eps=False, latent_type='uniform', delta_type='onedim', connect_mode='concat', batch_size=32, batch_per_gpu=16, return_atts=False, random_seed=1000, module_I_list=None, module_D_list=None, fmap_min=16, fmap_max=512, G_nf_scale=4, I_nf_scale=4, D_nf_scale=4, return_I_atts=False, dlatent_size=24, arch='resnet', topk_dims_to_show=20): # print('module_list:', module_list) train = EasyDict( run_func_name='training.training_loop_vc2.training_loop_vc2' ) # Options for training loop. D_global_size = 0 if not (module_list is None): module_list = _str_to_list(module_list) key_ls, size_ls, count_dlatent_size = split_module_names(module_list) for i, key in enumerate(key_ls): if key.startswith('D_global') or key.startswith('D_nocond_global'): D_global_size += size_ls[i] else: count_dlatent_size = dlatent_size if not (module_I_list is None): D_global_I_size = 0 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) for i, key in enumerate(key_I_ls): if key.startswith('D_global') or key.startswith('D_nocond_global'): D_global_I_size += size_I_ls[i] if not (module_D_list is None): D_global_D_size = 0 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) for i, key in enumerate(key_D_ls): if key.startswith('D_global') or key.startswith('D_nocond_global'): D_global_D_size += size_D_ls[i] if model_type == 'vc2_gan': # Standard VP-GAN G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.vc2_head', dlatent_size=count_dlatent_size, D_global_size=D_global_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_gan' elif model_type == 'vpex_gan': # VP-GAN extended. G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vpex_networks.vpex_net', dlatent_size=count_dlatent_size, D_global_size=D_global_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, return_atts=return_I_atts) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vpex_gan' elif model_type == 'vc2_gan_own_I': # Standard VP-GAN with own I G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.I_modular_vc2', dlatent_size=count_dlatent_I_size, D_global_size=D_global_I_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, module_I_list=module_I_list, I_nf_scale=I_nf_scale) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_gan_own_I' elif model_type == 'vc2_gan_own_ID': # Standard VP-GAN with own ID G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.I_modular_vc2', dlatent_size=count_dlatent_I_size, D_global_size=D_global_I_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, module_I_list=module_I_list, I_nf_scale=I_nf_scale) D = EasyDict(func_name='training.vc_networks2.D_modular_vc2', dlatent_size=count_dlatent_D_size, D_global_size=D_global_D_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, module_D_list=module_D_list, D_nf_scale=D_nf_scale) I_info = EasyDict() desc = 'vc2_gan_ownID' else: raise ValueError('Not supported model tyle: ' + model_type) 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. if model_type == 'vc2_gan': # Standard VP-GAN G_loss = EasyDict(func_name='training.loss_vc2.G_logistic_ns_vc2', D_global_size=D_global_size, C_lambda=C_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, delta_type=delta_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vpex_gan': # VP-GAN extended. G_loss = EasyDict(func_name='training.loss_vpex.G_logistic_ns_vpex', D_global_size=D_global_size, C_lambda=C_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, delta_type=delta_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vc2_gan_own_I' or model_type == 'vc2_gan_own_ID': # Standard VP-GAN with own I or D G_loss = EasyDict(func_name='training.loss_vc2.G_logistic_ns_vc2', D_global_size=D_global_size, C_lambda=C_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, delta_type=delta_type, own_I=True) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, 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': 1000} # Options for tflib.init_tf(). tf_config = { 'rnd.np_random_seed': random_seed } # 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 = batch_size sched.minibatch_gpu_base = batch_per_gpu D_loss.gamma = 10 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 desc += '-' + config_id # Configs A-E: Shrink networks to match original StyleGAN. I.fmap_base = 2 << I_fmap_base G.fmap_base = 2 << G_fmap_base D.fmap_base = 2 << D_fmap_base # Config E: Set gamma to 100 and override G & D architecture. # D_loss.gamma = 100 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, I_args=I, I_info_args=I_info, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss, use_info_gan=( model_type == 'vc2_info_gan2'), # Independent branch version use_vc_head=(model_type == 'vc2_gan' or model_type == 'vpex_gan' or model_type == 'vc2_gan_own_I' or model_type == 'vc2_gan_own_ID' or model_type == 'vc2_gan_byvae'), traversal_grid=True, return_atts=return_atts, return_I_atts=return_I_atts) n_continuous = 0 if not (module_list is None): for i, key in enumerate(key_ls): m_name = key.split('-')[0] if (m_name in LATENT_MODULES) and (not m_name == 'D_global'): n_continuous += size_ls[i] else: n_continuous = dlatent_size kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config, resume_pkl=resume_pkl, n_discrete=D_global_size, n_continuous=n_continuous, 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, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics, resume_pkl, I_fmap_base=8, G_fmap_base=8, D_fmap_base=9, fmap_decay=0.15, D_lambda=1, C_lambda=1, cls_alpha=0, n_samples_per=10, module_list=None, model_type='vc_gan2', epsilon_loss=3, random_eps=False, latent_type='uniform', delta_type='onedim', connect_mode='concat', batch_size=32, batch_per_gpu=16, return_atts=False, random_seed=1000, module_I_list=None, module_D_list=None, fmap_min=16, fmap_max=512, G_nf_scale=4, I_nf_scale=4, D_nf_scale=4, outlier_detector=False, gen_atts_in_D=False, no_atts_in_D=False, att_lambda=0, dlatent_size=24, arch='resnet', opt_reset_ls=None, norm_ord=2, n_dim_strict=0, drop_extra_torgb=False, latent_split_ls_for_std_gen=[5, 5, 5, 5], loose_rate=0.2, topk_dims_to_show=20, n_neg_samples=1, temperature=1., learning_rate=0.002, avg_mv_for_I=False, use_cascade=False, cascade_alt_freq_k=1, regW_lambda=1, network_snapshot_ticks=10): # print('module_list:', module_list) train = EasyDict( run_func_name='training.training_loop_vc2.training_loop_vc2' ) # Options for training loop. if opt_reset_ls is not None: opt_reset_ls = _str_to_list_of_int(opt_reset_ls) D_global_size = 0 if not (module_list is None): module_list = _str_to_list(module_list) key_ls, size_ls, count_dlatent_size = split_module_names(module_list) for i, key in enumerate(key_ls): if key.startswith('D_global') or key.startswith('D_nocond_global'): D_global_size += size_ls[i] else: count_dlatent_size = dlatent_size if not (module_I_list is None): D_global_I_size = 0 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) for i, key in enumerate(key_I_ls): if key.startswith('D_global') or key.startswith('D_nocond_global'): D_global_I_size += size_I_ls[i] if not (module_D_list is None): D_global_D_size = 0 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) for i, key in enumerate(key_D_ls): if key.startswith('D_global') or key.startswith('D_nocond_global'): D_global_D_size += size_D_ls[i] if model_type == 'vc2_info_gan': # G1 and G2 version InfoGAN G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. D = EasyDict(func_name='training.vc_networks2.D_info_modular_vc2', dlatent_size=count_dlatent_D_size, D_global_size=D_global_D_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, module_D_list=module_D_list, gen_atts_in_D=gen_atts_in_D, no_atts_in_D=no_atts_in_D, D_nf_scale=D_nf_scale) I = EasyDict() I_info = EasyDict() desc = 'vc2_info_gan_net' elif model_type == 'vc2_info_gan2': # Independent branch version InfoGAN G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.vc2_head_infogan2', dlatent_size=count_dlatent_size, D_global_size=D_global_size, fmap_min=fmap_min, fmap_max=fmap_max) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_info_gan2_net' elif model_type == 'vc2_gan': # Standard VP-GAN G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.vc2_head', dlatent_size=count_dlatent_size, D_global_size=D_global_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_gan' elif model_type == 'vc2_gan_byvae': # COMA-FAIN G = EasyDict( func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale, architecture=arch, drop_extra_torgb=drop_extra_torgb, latent_split_ls_for_std_gen=latent_split_ls_for_std_gen, ) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.vc2_head_byvae', dlatent_size=count_dlatent_size, D_global_size=D_global_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_gan_byvae' elif model_type == 'vc2_gan_byvae_simple': # COMA-FAIN-simple G = EasyDict( func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_simple_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale, architecture=arch, drop_extra_torgb=drop_extra_torgb, latent_split_ls_for_std_gen=latent_split_ls_for_std_gen, ) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.I_byvae_simple', dlatent_size=count_dlatent_size, D_global_size=D_global_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode) D = EasyDict(func_name='training.vc_networks2.D_stylegan2_simple', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_gan_byvae_simple' elif model_type == 'vc2_gan_style2_noI': # Just Style2-style GAN G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_stylegan2_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=dlatent_size, architecture=arch, dlatent_size=count_dlatent_size, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict() D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_gan_style2_noI' elif model_type == 'vc2_gan_own_I': # Standard VP-GAN with own I G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.I_modular_vc2', dlatent_size=count_dlatent_I_size, D_global_size=D_global_I_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, module_I_list=module_I_list, I_nf_scale=I_nf_scale) D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = 'vc2_gan_own_I' elif model_type == 'vc2_gan_own_ID': # Standard VP-GAN with own ID G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict(func_name='training.vc_networks2.I_modular_vc2', dlatent_size=count_dlatent_I_size, D_global_size=D_global_I_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, module_I_list=module_I_list, I_nf_scale=I_nf_scale) D = EasyDict(func_name='training.vc_networks2.D_modular_vc2', dlatent_size=count_dlatent_D_size, D_global_size=D_global_D_size, fmap_min=fmap_min, fmap_max=fmap_max, connect_mode=connect_mode, module_D_list=module_D_list, D_nf_scale=D_nf_scale) I_info = EasyDict() desc = 'vc2_gan_ownID' elif model_type == 'vc2_gan_noI' or model_type == 'vc2_traversal_contrastive' or \ model_type == 'gan_regW': # Just modular GAN or traversal contrastive or regW G = EasyDict(func_name='training.vc_networks2.G_main_vc2', synthesis_func='G_synthesis_modular_vc2', fmap_min=fmap_min, fmap_max=fmap_max, fmap_decay=fmap_decay, latent_size=count_dlatent_size, dlatent_size=count_dlatent_size, D_global_size=D_global_size, module_list=module_list, use_noise=True, return_atts=return_atts, G_nf_scale=G_nf_scale) # Options for generator network. I = EasyDict() D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2', fmap_min=fmap_min, fmap_max=fmap_max) # Options for discriminator network. I_info = EasyDict() desc = model_type else: raise ValueError('Not supported model tyle: ' + model_type) 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. if model_type == 'vc2_info_gan': # G1 and G2 version InfoGAN G_loss = EasyDict( func_name='training.loss_vc2.G_logistic_ns_vc2_info_gan', D_global_size=D_global_size, C_lambda=C_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, delta_type=delta_type, outlier_detector=outlier_detector, gen_atts_in_D=gen_atts_in_D, att_lambda=att_lambda) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2_info_gan', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vc2_info_gan2': # Independent branch version InfoGAN G_loss = EasyDict( func_name='training.loss_vc2.G_logistic_ns_vc2_info_gan2', D_global_size=D_global_size, C_lambda=C_lambda, latent_type=latent_type, norm_ord=norm_ord, n_dim_strict=n_dim_strict, loose_rate=loose_rate) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2_info_gan2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vc2_gan': # Standard VP-GAN G_loss = EasyDict(func_name='training.loss_vc2.G_logistic_ns_vc2', D_global_size=D_global_size, C_lambda=C_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, delta_type=delta_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vc2_gan_byvae' or model_type == 'vc2_gan_byvae_simple': # COMA-FAIN G_loss = EasyDict( func_name='training.loss_vc2.G_logistic_byvae_ns_vc2', D_global_size=D_global_size, C_lambda=C_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, use_cascade=use_cascade, delta_type=delta_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vc2_gan_own_I' or model_type == 'vc2_gan_own_ID': # Standard VP-GAN with own I or D G_loss = EasyDict(func_name='training.loss_vc2.G_logistic_ns_vc2', D_global_size=D_global_size, C_lambda=C_lambda, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, delta_type=delta_type, own_I=True) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vc2_gan_noI' or model_type == 'vc2_gan_style2_noI': # Just GANs (modular or StyleGAN2-style) G_loss = EasyDict( func_name='training.loss_vc2.G_logistic_ns', latent_type=latent_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'vc2_traversal_contrastive': # With perceptual distance as guide. G_loss = EasyDict( func_name= 'training.loss_vc2.G_logistic_ns_vc2_traversal_contrastive', D_global_size=D_global_size, C_lambda=C_lambda, n_neg_samples=n_neg_samples, temperature=temperature, epsilon=epsilon_loss, random_eps=random_eps, latent_type=latent_type, delta_type=delta_type) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, latent_type=latent_type) # Options for discriminator loss. elif model_type == 'gan_regW': G_loss = EasyDict( func_name='training.loss_vc2.G_logistic_ns_regW', latent_type=latent_type, regW_lambda=regW_lambda) # Options for generator loss. D_loss = EasyDict( func_name='training.loss_vc2.D_logistic_r1_vc2', D_global_size=D_global_size, 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': 1000} # Options for tflib.init_tf(). tf_config = { 'rnd.np_random_seed': random_seed } # 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.G_lrate_base = sched.D_lrate_base = learning_rate sched.minibatch_size_base = batch_size sched.minibatch_gpu_base = batch_per_gpu D_loss.gamma = 10 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 desc += '-' + config_id # Configs A-E: Shrink networks to match original StyleGAN. # I.fmap_base = 2 << 8 # G.fmap_base = 2 << 8 # D.fmap_base = 2 << 9 I.fmap_base = 2 << I_fmap_base G.fmap_base = 2 << G_fmap_base D.fmap_base = 2 << D_fmap_base # Config E: Set gamma to 100 and override G & D architecture. # D_loss.gamma = 100 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, I_args=I, I_info_args=I_info, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss, use_info_gan=( model_type == 'vc2_info_gan2'), # Independent branch version use_vc_head=(model_type == 'vc2_gan' or model_type == 'vc2_gan_own_I' or model_type == 'vc2_gan_own_ID' or model_type == 'vc2_gan_byvae' or model_type == 'vc2_gan_byvae_simple'), use_vc2_info_gan=(model_type == 'vc2_info_gan'), # G1 and G2 version use_perdis=(model_type == 'vc2_traversal_contrastive'), avg_mv_for_I=avg_mv_for_I, traversal_grid=True, return_atts=return_atts) n_continuous = 0 if not (module_list is None): for i, key in enumerate(key_ls): m_name = key.split('-')[0] if (m_name in LATENT_MODULES) and (not m_name == 'D_global'): n_continuous += size_ls[i] else: n_continuous = dlatent_size kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config, resume_pkl=resume_pkl, n_discrete=D_global_size, n_continuous=n_continuous, n_samples_per=n_samples_per, topk_dims_to_show=topk_dims_to_show, cascade_alt_freq_k=cascade_alt_freq_k, network_snapshot_ticks=network_snapshot_ticks) 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, mirror_augment, metrics, resume_G_pkl, n_batch=2, n_batch_per_gpu=1, D_global_size=0, C_global_size=10, model_type='hd_dis_model', latent_type='uniform', resume_pkl=None, n_samples_per=4, D_lambda=0, C_lambda=1, epsilon_in_loss=3, random_eps=True, M_lrmul=0.1, resolution_manual=1024, pretrained_type='with_stylegan2', traj_lambda=None, level_I_kimg=1000, use_level_training=False, resume_kimg=0, use_std_in_m=False, prior_latent_size=512, M_mapping_fmaps=512, hyperplane_lambda=1, hyperdir_lambda=1): train = EasyDict(run_func_name='training.training_loop_hd.training_loop_hd' ) # Options for training loop with pretrained HD. if model_type == 'hd_hyperplane': M = EasyDict( func_name='training.hd_networks.net_M_hyperplane', C_global_size=C_global_size, D_global_size=D_global_size, latent_size=prior_latent_size, mapping_lrmul=M_lrmul, use_std_in_m=use_std_in_m) # Options for dismapper network. I = EasyDict( func_name='training.hd_networks.net_I', C_global_size=C_global_size, D_global_size=D_global_size) # Options for recognizor network. else: M = EasyDict( func_name='training.hd_networks.net_M', C_global_size=C_global_size, D_global_size=D_global_size, latent_size=prior_latent_size, mapping_fmaps=M_mapping_fmaps, mapping_lrmul=M_lrmul, use_std_in_m=use_std_in_m) # Options for dismapper network. I = EasyDict( func_name='training.hd_networks.net_I', C_global_size=C_global_size, D_global_size=D_global_size) # Options for recognizor network. if model_type == 'hd_dis_model_with_cls': I_info = EasyDict(func_name='training.hd_networks.net_I_info', C_global_size=C_global_size, D_global_size=D_global_size) else: I_info = EasyDict() I_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. if model_type == 'hd_hyperplane': I_loss = EasyDict( func_name='training.loss_hd.IandM_hyperplane_loss', latent_type=latent_type, D_global_size=D_global_size, C_global_size=C_global_size, D_lambda=D_lambda, C_lambda=C_lambda, epsilon=epsilon_in_loss, random_eps=random_eps, traj_lambda=traj_lambda, resolution_manual=resolution_manual, use_std_in_m=use_std_in_m, model_type=model_type, hyperplane_lambda=hyperplane_lambda, prior_latent_size=prior_latent_size, hyperdir_lambda=hyperdir_lambda) # Options for discriminator loss. else: I_loss = EasyDict(func_name='training.loss_hd.IandM_loss', latent_type=latent_type, D_global_size=D_global_size, C_global_size=C_global_size, D_lambda=D_lambda, C_lambda=C_lambda, epsilon=epsilon_in_loss, random_eps=random_eps, traj_lambda=traj_lambda, resolution_manual=resolution_manual, use_std_in_m=use_std_in_m, model_type=model_type, hyperplane_lambda=hyperplane_lambda, prior_latent_size=prior_latent_size ) # 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.data_dir = data_dir train.total_kimg = total_kimg train.mirror_augment = mirror_augment train.image_snapshot_ticks = train.network_snapshot_ticks = 10 sched.I_lrate_base = 0.002 sched.minibatch_size_base = n_batch sched.minibatch_gpu_base = n_batch_per_gpu metrics = [metric_defaults[x] for x in metrics] desc = 'hd_disentanglement' 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': I.fmap_base = 8 << 10 sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(I_args=I, M_args=M, I_opt_args=I_opt, I_loss_args=I_loss, resume_G_pkl=resume_G_pkl) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, use_hd_with_cls=(model_type == 'hd_dis_model_with_cls'), use_hyperplane=(model_type == 'hd_hyperplane'), metric_arg_list=metrics, tf_config=tf_config, resume_pkl=resume_pkl, n_discrete=D_global_size, n_continuous=C_global_size, n_samples_per=n_samples_per, resolution_manual=resolution_manual, pretrained_type=pretrained_type, level_I_kimg=level_I_kimg, use_level_training=use_level_training, resume_kimg=resume_kimg, use_std_in_m=use_std_in_m, prior_latent_size=prior_latent_size) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)