def run(data_dir='./../datasets',
result_dir='./../results',
num_gpus=1,
total_kimg=1000,
mirror_augment=True):
train = EasyDict(
run_func_name='classifier_vgg.training_loop-vgg_rotation.training_loop'
)
classifier = EasyDict(
func_name='classifier_vgg.network_classifier-vgg.classifier_vgg')
classifier_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
classifier_loss = EasyDict(func_name='classifier_vgg.loss.cross_entropy')
sched = EasyDict()
sc = dnnlib.SubmitConfig()
tf_config = {'rnd.np_random_seed': 1000}
# train.resume_pkl = './results/00254-classifier-single_class_model/network-snapshot-001000.pkl'
# train.resume_kimg = 1000.0
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 8
current_label = 'rotation'
dataset = 'classifier_oversample_' + current_label
print(dataset)
train.data_dir = data_dir
if not os.path.exists(train.data_dir):
print('Error: dataset root directory does not exist.')
sys.exit(1)
train.total_kimg = total_kimg
train.mirror_augment = True
metrics = [
EasyDict(name='acc_' + current_label,
func_name='metrics.accuracy_single_class.ACC',
test_dataset=dataset + '_test',
num_images=5000,
minibatch_per_gpu=8,
test_data_dir=data_dir)
]
desc = 'classifier-single_class_' + current_label
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(classifier_args=classifier,
classifier_opt_args=classifier_opt,
classifier_loss_args=classifier_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
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, config_id, num_gpus, total_kimg, gamma,
mirror_augment, metrics):
train = EasyDict(
run_func_name='training.training_loop.rotation.v5_int_reg.training_loop'
)
G = EasyDict(func_name='training.networks.rotation.v5_int_reg.G_main')
D = EasyDict(func_name='training.networks.rotation.v5_int_reg.D_stylegan2')
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
G_loss = EasyDict(
func_name='training.loss.rotation.v5_int_reg.G_logistic_ns_pathreg')
D_loss = EasyDict(
func_name='training.loss.rotation.v5_int_reg.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 = 32
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'rotation-v5-int-reg_256'
G_loss.int_reg_clip = 5.0
G_loss.rotation_step_size = 0.08 / 2
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
assert config_id in _valid_configs
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(data_dir='./../datasets',
result_dir='./../results',
num_gpus=2,
total_kimg=1000):
train = EasyDict(
run_func_name='classifier.training_loop-single_class.training_loop')
classifier = EasyDict(
func_name='classifier.network_classifier-new_label.classifier')
classifier_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
classifier_loss = EasyDict(func_name='classifier.loss.euclidean')
sched = EasyDict()
sc = dnnlib.SubmitConfig()
tf_config = {'rnd.np_random_seed': 1000}
dataset = 'classifier_oversample_rotation_v7'
sched.minibatch_size_base = 8
sched.minibatch_gpu_base = 4
train.data_dir = data_dir
train.total_kimg = total_kimg
desc = 'classifier_rotation_v7'
metrics = [
EasyDict(
name='ACCv7',
func_name='metrics.accuracy_v7.ACCv7',
num_images=5000,
minibatch_per_gpu=8,
test_dataset='./../datasets/classifier_oversample_rotation_v7_test'
)
]
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(classifier_args=classifier,
classifier_opt_args=classifier_opt,
classifier_loss_args=classifier_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
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_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.conditional.v5_baseline.training_loop') # Options for training loop.
G = EasyDict(func_name='training.networks.conditional.baseline.G_main') # Options for generator network.
D = EasyDict(func_name='training.networks.conditional.baseline.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.conditional.label_dropout.G_logistic_ns_pathreg') # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.conditional.label_dropout.D_logistic_r1') # 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.resume_pkl = './../results/00326-conditional_label_dropout_25/network-snapshot-000887.pkl'
# train.resume_kimg = 887.0
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 = 'conditional_label_dropout_25'
G_loss.label_dropout_prob = 0.5
D_loss.label_dropout_prob = 0.5
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
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,
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,
config_id,
num_gpus,
total_kimg,
gamma,
mirror_augment,
metrics,
resume_pkl,
D_global_size=3,
C_global_size=0, # Global C_latents.
sb_C_global_size=4,
C_local_hfeat_size=0, # Local heatmap*features learned C_latents.
C_local_heat_size=0, # Local heatmap learned C_latents.
n_samples_per=10,
module_list=None,
single_const=True,
model_type='spatial_biased'):
# print('module_list:', module_list)
train = EasyDict(
run_func_name='training.training_loop_dsp.training_loop_dsp'
) # Options for training loop.
if model_type == 'spatial_biased':
G = EasyDict(
func_name=
'training.spatial_biased_networks.G_main_spatial_biased_dsp',
mapping_fmaps=128,
fmap_max=128,
latent_size=D_global_size + sb_C_global_size,
dlatent_size=D_global_size + sb_C_global_size,
D_global_size=D_global_size,
sb_C_global_size=sb_C_global_size
) # Options for generator network.
desc = 'spatial_biased_net'
elif model_type == 'sb_general':
G = EasyDict(
func_name=
'training.spatial_biased_networks.G_main_spatial_biased_dsp',
synthesis_func='G_synthesis_sb_general_dsp',
mapping_fmaps=128,
fmap_max=128,
latent_size=D_global_size + C_global_size + sb_C_global_size +
C_local_hfeat_size + C_local_heat_size,
dlatent_size=D_global_size + C_global_size + sb_C_global_size +
C_local_hfeat_size + C_local_heat_size,
D_global_size=D_global_size,
C_global_size=C_global_size,
sb_C_global_size=sb_C_global_size,
C_local_hfeat_size=C_local_hfeat_size,
C_local_heat_size=C_local_heat_size,
use_noise=False) # Options for generator network.
desc = 'sb_general_net'
elif model_type == 'sb_modular':
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'):
D_global_size = size_ls[i]
break
print('D_global_size:', D_global_size)
G = EasyDict(
func_name=
'training.spatial_biased_networks.G_main_spatial_biased_dsp',
synthesis_func='G_synthesis_sb_modular',
mapping_fmaps=128,
fmap_max=128,
latent_size=count_dlatent_size,
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
module_list=module_list,
single_const=single_const,
use_noise=False) # Options for generator network.
desc = 'sb_modular_net'
elif model_type == 'sb_singlelayer_modi':
G = EasyDict(func_name='training.simple_networks.G_main_simple_dsp',
synthesis_func='G_synthesis_sb_singlelayer_modi_dsp',
mapping_fmaps=128,
fmap_max=128,
latent_size=D_global_size + sb_C_global_size,
dlatent_size=D_global_size + sb_C_global_size,
D_global_size=D_global_size,
sb_C_global_size=sb_C_global_size
) # Options for generator network.
desc = 'sb_singlelayer_net'
elif model_type == 'stylegan2':
G = EasyDict(
func_name=
'training.spatial_biased_networks.G_main_spatial_biased_dsp',
dlatent_avg_beta=None,
mapping_fmaps=128,
fmap_max=128,
latent_size=12,
D_global_size=D_global_size,
sb_C_global_size=sb_C_global_size
) # Options for generator network.
desc = 'stylegan2_net'
elif model_type == 'simple':
G = EasyDict(func_name='training.simple_networks.G_main_simple_dsp',
latent_size=D_global_size + sb_C_global_size,
dlatent_size=D_global_size + sb_C_global_size,
D_global_size=D_global_size,
sb_C_global_size=sb_C_global_size
) # Options for generator network.
else:
raise ValueError('Not supported model tyle: ' + model_type)
if model_type == 'simple':
D = EasyDict(func_name='training.simple_networks.D_simple_dsp'
) # Options for discriminator network.
else:
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2',
fmap_max=128) # Options for discriminator network.
# D = EasyDict(func_name='training.spatial_biased_networks.D_with_discrete_dsp', fmap_max=128) # 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_dsp',
D_global_size=D_global_size) # Options for generator loss.
D_loss = EasyDict(
func_name='training.loss.D_logistic_r1_dsp',
D_global_size=D_global_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.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 += '-' + 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
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)
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,
traversal_grid=True)
if model_type == 'sb_modular':
n_continuous = 0
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 = C_global_size + sb_C_global_size + \
C_local_hfeat_size + C_local_heat_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)
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, gamma, mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.training_loop_mirror_v6_remove_half_fl_fr.training_loop')
G = EasyDict(func_name='training.networks.networks_stylegan2.G_main')
D = EasyDict(func_name='training.networks.networks_stylegan2_discriminator_new_rotation.D_stylegan2_new_rotaion') # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
G_loss = EasyDict(func_name='training.loss.loss_G_new_rotation_squared_euclidean_10_interpolate_50_percent_uniform_dist_int_penalty.G_logistic_ns_pathreg')
D_loss = EasyDict(func_name='training.loss.loss_D_logistic_r1_new_rotation_euclidean_square.D_logistic_r1_new_rotation')
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 = 32
sched.minibatch_gpu_base = 4
# train.resume_pkl = './results/00200-stylegan2-car_labels_v7_oversample_filter-2gpu-config-f-squared_euclidean_10_interpolate_50_percent_int_reg-256/network-snapshot-000887.pkl'
# train.resume_kimg = 887.2
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
G.style_mixing_prob = None
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
desc += '-squared_euclidean_10_interpolate_50_percent_int_reg_remove_half_fl_fr_no_noise_square'
desc += '-256'
# 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, 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)
def run(
dataset,
data_dir,
result_dir,
config_id,
num_gpus,
total_kimg,
gamma,
mirror_augment,
metrics,
resume_pkl=None,
resume_kimg=None,
):
train = EasyDict(
run_func_name="training.training_loop.training_loop",
# training resume options:
resume_pkl=
resume_pkl, # Network pickle to resume training from, None = train from scratch.
resume_kimg=
resume_kimg, # Assumed training progress at the beginning. Affects reporting and training schedule.
) # 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 != "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, config_id, num_gpus, total_kimg, gamma,
mirror_augment, metrics, dlatent_size, lr, batch_size, decay_step,
decay_rate, stair, tick_kimg):
train = EasyDict(run_func_name='training.vae_training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.vae_dcgan.Decoder_main'
) # Options for generator network.
D = EasyDict(func_name='training.vae_dcgan.Encoder'
) # 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.vae_loss') # 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.batch_size = batch_size
sched.lr = lr
sched.decay_step = decay_step
sched.decay_rate = decay_rate
sched.stair = stair
sched.tick_kimg = tick_kimg
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'vae_dcgan'
G.dlatent_size = dlatent_size
D.dlatent_size = dlatent_size
G.num_units = D.num_units = 1024
G.act = D.act = 'relu'
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
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,
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_run_id=None):
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().
if resume_run_id is not None:
# Resume from the ID of the results directory given
ids = sorted(get_valid_runids(result_dir))
if resume_run_id == 'recent':
resume_run_id = ids[-1][0]
else:
try:
resume_run_id = int(resume_run_id)
except ValueError:
raise RuntimeError(
'--resume argument is invalid (must be number, or "recent"): {}'
.format(resume_run_id))
try:
rundir_name = next(x[1] for x in ids if x[0] == resume_run_id)
except StopIteration:
raise RuntimeError(
'Could not find results directory with run ID {} (options: {})'
.format(resume_run_id, [x[0] for x in ids]))
# Find kimg & pkl file
rundir = os.path.join(result_dir, rundir_name)
pkls = [
name for name in os.listdir(rundir)
if name.startswith('network-snapshot-') and name.endswith('.pkl')
]
kimgs = sorted([(int(
pkl.replace('network-snapshot-', '').replace('.pkl', '')), pkl)
for pkl in pkls],
key=lambda x: x[0])
if len(kimgs) == 0:
raise RuntimeError(
'No network-snapshot-[0-9].pkl files found in {}'.format(
rundir))
max_kimg = kimgs[-1][0]
pkl_name = kimgs[-1][1]
# Get wall clock time
logfilepath = os.path.join(rundir, 'log.txt')
with open(logfilepath, 'r') as f:
logfile = f.read()
for line in logfile.splitlines():
if 'kimg {}'.format(max_kimg) in line:
if 'time ' not in line:
raise RuntimeError(
'Invalid log file: {}'.format(logfilepath))
line = line.split('time ')[1]
if 'sec/tick' not in line:
raise RuntimeError(
'Invalid log file: {}'.format(logfilepath))
line = line.split('sec/tick')[0].strip()
# Parse d h m s, etc.
total_seconds_formatted = line
total_seconds = 0
if 'd' in line:
arr = line.split('d')
days = int(arr[0].strip())
total_seconds += days * 24 * 60 * 60
line = arr[1]
if 'h' in line:
arr = line.split('h')
hours = int(arr[0].strip())
total_seconds += hours * 60 * 60
line = arr[1]
if 'm' in line:
arr = line.split('m')
mins = int(arr[0].strip())
total_seconds += mins * 60
line = arr[1]
if 's' in line:
arr = line.split('s')
secs = int(arr[0].strip())
total_seconds += secs
line = arr[1]
break
# Set args for training
train.resume_pkl = os.path.join(rundir, pkl_name)
train.resume_kimg = max_kimg
train.resume_time = total_seconds
print('Resuming from run {}: kimg {}, time {}'.format(
rundir_name, max_kimg, total_seconds_formatted))
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 1
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,
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(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)
def run(dataset, data_dir, result_dir, num_gpus, total_kimg, mirror_augment,
metrics, resume, resume_with_new_nets, disable_style_mod,
disable_cond_mod):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.co_mod_gan.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.co_mod_gan.D_co_mod_gan'
) # 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_masked_logistic_ns_l1'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_masked_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 = 'co-mod-gan'
desc += '-' + os.path.basename(dataset)
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
if resume is not None:
resume_kimg = int(
os.path.basename(resume).replace('.pkl', '').split('-')[-1])
else:
resume_kimg = 0
if disable_style_mod:
G.style_mod = False
if disable_cond_mod:
G.cond_mod = False
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.update(resume_pkl=resume,
resume_kimg=resume_kimg,
resume_with_new_nets=resume_with_new_nets)
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,
fmap_decay=0.15,
D_lambda=1,
C_lambda=1,
MI_lambda=1,
cls_alpha=0,
n_samples_per=10,
module_list=None,
single_const=True,
model_type='spatial_biased',
phi_blurry=0.5,
latent_type='uniform'):
train = EasyDict(
run_func_name='training.training_loop_vid.training_loop_vid'
) # Options for training loop.
D_global_size = 0
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'):
D_global_size += size_ls[i]
break
print('D_global_size:', D_global_size)
print('key_ls:', key_ls)
print('size_ls:', size_ls)
print('count_dlatent_size:', count_dlatent_size)
if model_type == 'vid_model':
G = EasyDict(func_name='training.vid_networks.G_main_vid',
synthesis_func='G_synthesis_vid_modular',
fmap_min=16,
fmap_max=512,
fmap_decay=fmap_decay,
latent_size=count_dlatent_size,
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
module_list=module_list,
single_const=single_const,
use_noise=True) # Options for generator network.
I = EasyDict(func_name='training.vid_networks.vid_head',
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
fmap_max=512)
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2',
fmap_max=512) # Options for discriminator network.
I_info = EasyDict()
desc = model_type
elif model_type == 'vid_with_cls':
G = EasyDict(func_name='training.vid_networks.G_main_vid',
synthesis_func='G_synthesis_vid_modular',
fmap_min=16,
fmap_max=512,
fmap_decay=fmap_decay,
latent_size=count_dlatent_size,
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
module_list=module_list,
single_const=single_const,
use_noise=True) # Options for generator network.
I = EasyDict(func_name='training.vid_networks.vid_head',
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
fmap_max=512)
I_info = EasyDict(
func_name='training.info_gan_networks.info_gan_head_cls',
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
fmap_decay=fmap_decay,
fmap_min=16,
fmap_max=512)
D = EasyDict(
func_name='training.info_gan_networks.D_info_gan_stylegan2',
fmap_max=512) # Options for discriminator network.
desc = model_type
elif model_type == 'vid_naive_cluster_model':
G = EasyDict(func_name='training.vid_networks.G_main_vid',
synthesis_func='G_synthesis_vid_modular',
fmap_min=16,
fmap_max=512,
fmap_decay=fmap_decay,
latent_size=count_dlatent_size,
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
module_list=module_list,
single_const=single_const,
use_noise=True) # Options for generator network.
I = EasyDict(func_name='training.vid_networks.vid_naive_cluster_head',
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
fmap_max=512) # Options for estimator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2',
fmap_max=512) # Options for discriminator network.
I_info = EasyDict()
desc = model_type
elif model_type == 'vid_blurry_model':
G = EasyDict(func_name='training.vid_networks.G_main_vid',
synthesis_func='G_synthesis_vid_modular',
fmap_min=16,
fmap_max=512,
fmap_decay=fmap_decay,
latent_size=count_dlatent_size,
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
module_list=module_list,
single_const=single_const,
use_noise=True) # Options for generator network.
I = EasyDict(func_name='training.vid_networks.vid_naive_cluster_head',
dlatent_size=count_dlatent_size,
D_global_size=D_global_size,
fmap_max=512) # Options for estimator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2',
fmap_max=512) # 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.
I_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
if model_type == 'vid_model':
G_loss = EasyDict(
func_name='training.loss_vid.G_logistic_ns_vid',
D_global_size=D_global_size,
C_lambda=C_lambda,
latent_type=latent_type) # Options for generator loss.
D_loss = EasyDict(
func_name='training.loss_vid.D_logistic_r1_vid',
D_global_size=D_global_size,
latent_type=latent_type) # Options for discriminator loss.
I_loss = EasyDict(func_name='training.loss_vid.I_vid',
D_global_size=D_global_size,
latent_type=latent_type,
C_lambda=C_lambda,
MI_lambda=MI_lambda) # Options for estimator loss.
elif model_type == 'vid_with_cls':
G_loss = EasyDict(
func_name='training.loss_vid.G_logistic_ns_vid',
D_global_size=D_global_size,
C_lambda=C_lambda,
cls_alpha=cls_alpha,
latent_type=latent_type) # Options for generator loss.
D_loss = EasyDict(
func_name='training.loss_vid.D_logistic_r1_info_gan_vid',
D_global_size=D_global_size,
latent_type=latent_type) # Options for discriminator loss.
I_loss = EasyDict(func_name='training.loss_vid.I_vid',
D_global_size=D_global_size,
latent_type=latent_type,
C_lambda=C_lambda,
MI_lambda=MI_lambda) # Options for estimator loss.
elif model_type == 'vid_naive_cluster_model':
G_loss = EasyDict(
func_name='training.loss_vid.G_logistic_ns_vid_naive_cluster',
D_global_size=D_global_size,
C_lambda=C_lambda,
latent_type=latent_type) # Options for generator loss.
D_loss = EasyDict(
func_name='training.loss_vid.D_logistic_r1_vid',
D_global_size=D_global_size,
latent_type=latent_type) # Options for discriminator loss.
I_loss = EasyDict() # Options for estimator loss.
I_opt = EasyDict()
elif model_type == 'vid_blurry_model':
G_loss = EasyDict(
func_name='training.loss_vid.G_logistic_ns_vid_naive_cluster',
D_global_size=D_global_size,
C_lambda=C_lambda,
latent_type=latent_type) # Options for generator loss.
D_loss = EasyDict(
func_name='training.loss_vid.D_logistic_r1_vid',
D_global_size=D_global_size,
latent_type=latent_type) # Options for discriminator loss.
I_loss = EasyDict(func_name='training.loss_vid.I_vid_blurry',
D_global_size=D_global_size,
latent_type=latent_type,
C_lambda=C_lambda,
MI_lambda=MI_lambda,
phi=phi_blurry) # Options for estimator loss.
else:
raise ValueError('Not supported loss tyle: ' + 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': 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 = sched.I_lrate_base = 0.002
sched.minibatch_size_base = 16
sched.minibatch_gpu_base = 8
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
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 = G.fmap_base = D.fmap_base = 8 << 10
I.fmap_base = G.fmap_base = D.fmap_base = 2 << 8
# 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)
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,
I_opt_args=I_opt,
G_loss_args=G_loss,
D_loss_args=D_loss,
I_loss_args=I_loss,
use_vid_head=(model_type == 'vid_model'),
use_vid_head_with_cls=(model_type == 'vid_with_cls'),
use_vid_naive_cluster=(model_type == 'vid_naive_cluster_model'),
use_vid_blurry=(model_type == 'vid_blurry_model'),
traversal_grid=True)
n_continuous = 0
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_discrete=D_global_size,
n_continuous=n_continuous,
n_samples_per=n_samples_per,
C_lambda=C_lambda,
MI_lambda=MI_lambda)
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().
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = 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,
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)