def run(data, train_dir, config, d_aug, diffaug_policy, cond, ops, mirror, mirror_v, \
kimg, batch_size, lrate, resume, resume_kimg, num_gpus, ema_kimg, gamma, freezeD):
# training functions
if d_aug: # https://github.com/mit-han-lab/data-efficient-gans
train = EasyDict(
run_func_name='training.training_loop_diffaug.training_loop'
) # Options for training loop (Diff Augment method)
loss_args = EasyDict(
func_name='training.loss_diffaug.ns_DiffAugment_r1',
policy=diffaug_policy) # Options for loss (Diff Augment method)
else: # original nvidia
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop (original from NVidia)
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
# network functions
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='1080p',
layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
G.impl = D.impl = ops
# dataset (tfrecords) - get or create
tfr_files = file_list(os.path.dirname(data), 'tfr')
tfr_files = [
f for f in tfr_files if basename(data) == basename(f).split('-')[0]
]
if len(tfr_files) == 0 or os.stat(tfr_files[0]).st_size == 0:
tfr_file, total_samples = create_from_image_folders(
data) if cond is True else create_from_images(data)
else:
tfr_file = tfr_files[0]
dataset_args = EasyDict(tfrecord=tfr_file)
# resolutions
with tf.Graph().as_default(), tflib.create_session().as_default(): # pylint: disable=not-context-manager
dataset_obj = dataset.load_dataset(
**dataset_args) # loading the data to see what comes out
resolution = dataset_obj.resolution
init_res = dataset_obj.init_res
res_log2 = dataset_obj.res_log2
dataset_obj.close()
dataset_obj = None
if list(init_res) == [4, 4]:
desc = '%s-%d' % (basename(data), resolution)
else:
print(' custom init resolution', init_res)
desc = basename(tfr_file)
G.init_res = D.init_res = list(init_res)
train.savenames = [desc.replace(basename(data), 'snapshot'), desc]
desc += '-%s' % config
# training schedule
train.total_kimg = kimg
train.image_snapshot_ticks = 1 * num_gpus if kimg <= 1000 else 4 * num_gpus
train.network_snapshot_ticks = 5
train.mirror_augment = mirror
train.mirror_augment_v = mirror_v
sched.tick_kimg_base = 2 if train.total_kimg < 2000 else 4
# learning rate
if config == 'e':
sched.G_lrate_base = 0.001
sched.G_lrate_dict = {0: 0.001, 1: 0.0007, 2: 0.0005, 3: 0.0003}
sched.lrate_step = 1500 # period for stepping to next lrate, in kimg
if config == 'f':
sched.G_lrate_base = lrate # 0.001 for big datasets, 0.0003 for few-shot
sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway
# batch size (for 16gb memory GPU)
sched.minibatch_gpu_base = 4096 // resolution if batch_size is None else batch_size
print(' Batch size', sched.minibatch_gpu_base)
sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base
sc.num_gpus = num_gpus
if config == 'e':
G.fmap_base = D.fmap_base = 8 << 10
if d_aug: loss_args.gamma = 100 if gamma is None else gamma
else: D_loss.gamma = 100 if gamma is None else gamma
elif config == 'f':
G.fmap_base = D.fmap_base = 16 << 10
else:
print(' Only configs E and F are implemented')
exit()
if cond:
desc += '-cond'
dataset_args.max_label_size = 'full' # conditioned on full label
if freezeD:
D.freezeD = True
train.resume_with_new_nets = True
if d_aug:
desc += '-daug'
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
tf_config=tf_config)
kwargs.update(resume_pkl=resume,
resume_kimg=resume_kimg,
resume_with_new_nets=True)
if ema_kimg is not None:
kwargs.update(G_ema_kimg=ema_kimg)
if d_aug:
kwargs.update(loss_args=loss_args)
else:
kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = train_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(dataset,
data_dir,
result_dir,
config_id,
num_gpus,
total_kimg,
gamma,
mirror_augment,
metrics,
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, resolution, result_dir, DiffAugment, num_gpus, batch_size,
total_kimg, ema_kimg, num_samples, gamma, fmap_base, fmap_max,
latent_size, mirror_augment, impl, metrics, resume, resume_kimg,
num_repeats, eval):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
loss_args = EasyDict(
func_name='training.loss.ns_r1_DiffAugment') # Options for loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='4k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
metrics = [metric_defaults[x] for x in metrics]
metric_args = EasyDict(num_repeats=num_repeats)
desc = 'DiffAugment-stylegan2' if DiffAugment else 'stylegan2'
dataset_args = EasyDict(tfrecord_dir=dataset, resolution=resolution)
desc += '-' + os.path.basename(dataset)
if resolution is not None:
desc += '-{}'.format(resolution)
if num_samples is not None:
dataset_args.num_samples = num_samples
desc += '-{}samples'.format(num_samples)
if batch_size is not None:
desc += '-batch{}'.format(batch_size)
else:
batch_size = 32
assert batch_size % num_gpus == 0
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
sched.minibatch_size_base = batch_size
sched.minibatch_gpu_base = batch_size // num_gpus
G.impl = D.impl = impl
if fmap_base is not None:
G.fmap_base = D.fmap_base = fmap_base
desc += '-fmap{}'.format(fmap_base)
if fmap_max is not None:
G.fmap_max = D.fmap_max = fmap_max
desc += '-fmax{}'.format(fmap_max)
if latent_size is not None:
G.latent_size = G.mapping_fmaps = G.dlatent_size = latent_size
desc += '-latent{}'.format(latent_size)
if gamma is not None:
loss_args.gamma = gamma
desc += '-gamma{}'.format(gamma)
if DiffAugment:
loss_args.policy = DiffAugment
desc += '-' + DiffAugment.replace(',', '-')
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
loss_args=loss_args)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.update(resume_pkl=resume,
resume_kimg=resume_kimg,
resume_with_new_nets=True)
kwargs.update(metric_args=metric_args)
if ema_kimg is not None:
kwargs.update(G_ema_kimg=ema_kimg)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def main():
parser = argparse.ArgumentParser(
description='''StyleGAN2 generator.
Run 'python %(prog)s <subcommand> --help' for subcommand help.''',
epilog=_examples,
formatter_class=argparse.RawDescriptionHelpFormatter)
subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
parser_generate_images = subparsers.add_parser('generate-images',
help='Generate images')
parser_generate_images.add_argument('--network',
help='Network pickle filename',
dest='network_pkl',
required=True)
parser_generate_images.add_argument('--seeds',
type=_parse_num_range,
help='List of random seeds',
required=True)
parser_generate_images.add_argument(
'--result-dir',
help='Root directory for run results (default: %(default)s)',
default='results',
metavar='DIR')
parser_generate_images.add_argument(
'--create_new_G',
help='If create a new G for projection.',
default=False,
type=_str_to_bool)
parser_generate_images.add_argument(
'--new_func_name',
help='new G func name if create new G',
default='training.vc_networks2.G_main_vc2')
parser_style_mixing_example = subparsers.add_parser(
'style-mixing-example', help='Generate style mixing video')
parser_style_mixing_example.add_argument('--network',
help='Network pickle filename',
dest='network_pkl',
required=True)
parser_style_mixing_example.add_argument(
'--row-seeds',
type=_parse_num_range,
help='Random seeds to use for image rows',
required=True)
parser_style_mixing_example.add_argument(
'--col-seeds',
type=_parse_num_range,
help='Random seeds to use for image columns',
required=True)
parser_style_mixing_example.add_argument(
'--col-styles',
type=_parse_num_range,
help='Style layer range (default: %(default)s)',
default='0-6')
parser_style_mixing_example.add_argument(
'--truncation-psi',
type=float,
help='Truncation psi (default: %(default)s)',
default=0.5)
parser_style_mixing_example.add_argument(
'--result-dir',
help='Root directory for run results (default: %(default)s)',
default='results',
metavar='DIR')
parser_generate_travs = subparsers.add_parser('generate-traversals',
help='Generate traversals')
parser_generate_travs.add_argument('--network',
help='Network pickle filename',
dest='network_pkl',
required=True)
parser_generate_travs.add_argument('--seeds',
type=_parse_num_range,
help='List of random seeds',
required=True)
parser_generate_travs.add_argument(
'--result-dir',
help='Root directory for run results (default: %(default)s)',
default='results',
metavar='DIR')
parser_generate_travs.add_argument('--tpl_metric',
help='TPL to use',
default='tpl',
type=str)
parser_generate_travs.add_argument('--n_samples_per',
help='N samplers per row',
default=7,
type=int)
parser_generate_travs.add_argument('--topk_dims_to_show',
help='Top k dims to show',
default=-1,
type=int)
parser_generate_travs.add_argument('--return_atts',
help='If save atts.',
default=False,
type=_str_to_bool)
parser_generate_travs.add_argument('--bound',
help='Traversal bound',
default=2,
type=float)
parser_generate_domain_shift = subparsers.add_parser(
'generate-domain-shift', help='Generate traversals')
parser_generate_domain_shift.add_argument('--network',
help='Network pickle filename',
dest='network_pkl',
required=True)
parser_generate_domain_shift.add_argument('--seeds',
type=_parse_num_range,
help='List of random seeds',
required=True)
parser_generate_domain_shift.add_argument(
'--result-dir',
help='Root directory for run results (default: %(default)s)',
default='results',
metavar='DIR')
parser_generate_domain_shift.add_argument('--domain_dim',
help='The dim denoting domain',
default=0,
type=int)
parser_generate_gifs = subparsers.add_parser('generate-gifs',
help='Generate gifs')
parser_generate_gifs.add_argument('--network',
help='Network pickle filename',
dest='network_pkl',
required=True)
parser_generate_gifs.add_argument('--exist_imgs_dir',
help='Dir for used input images',
default='inputs',
metavar='EXIST')
parser_generate_gifs.add_argument(
'--result-dir',
help='Root directory for run results (default: %(default)s)',
default='results',
metavar='DIR')
parser_generate_gifs.add_argument('--used_imgs_ls',
help='Image names to use',
default='[img1.png, img2.png]',
type=_str_to_list)
parser_generate_gifs.add_argument('--used_semantics_ls',
help='Semantics to use',
default='[azimuth, haircolor]',
type=_str_to_list)
parser_generate_gifs.add_argument(
'--attr2idx_dict',
help='Attr names to attr idx in latent codes',
default='{azimuth: 10, haircolor: 17, smile: 6}',
type=_str_to_attr2idx)
parser_generate_gifs.add_argument('--create_new_G',
help='If create a new G for projection.',
default=False,
type=_str_to_bool)
parser_generate_gifs.add_argument(
'--new_func_name',
help='new G func name if create new G',
default='training.vc_networks2.G_main_vc2')
args = parser.parse_args()
kwargs = vars(args)
subcmd = kwargs.pop('command')
if subcmd is None:
print('Error: missing subcommand. Re-run with --help for usage.')
sys.exit(1)
sc = dnnlib.SubmitConfig()
sc.num_gpus = 1
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
sc.run_dir_root = kwargs.pop('result_dir')
sc.run_desc = subcmd
func_name_map = {
'generate-images': 'run_generator_vc2.generate_images',
'style-mixing-example': 'run_generator_vc2.style_mixing_example',
'generate-traversals': 'run_generator_vc2.generate_traversals',
'generate-domain-shift': 'run_generator_vc2.generate_domain_shift',
'generate-gifs': 'run_generator_vc2.generate_gifs'
}
dnnlib.submit_run(sc, func_name_map[subcmd], **kwargs)
def main():
parser = argparse.ArgumentParser(
description='''StyleGAN2 generator.
Run 'python %(prog)s <subcommand> --help' for subcommand help.''',
epilog=_examples,
formatter_class=argparse.RawDescriptionHelpFormatter
)
subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
parser_generate_images = subparsers.add_parser('generate-images', help='Generate images')
parser_generate_images.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
parser_generate_images.add_argument('--seeds', type=_parse_num_range, help='List of random seeds', required=True)
parser_generate_images.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
parser_generate_images.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
parser_generate_ws = subparsers.add_parser('generate-ws', help='Generate Ws')
parser_generate_ws.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
parser_generate_ws.add_argument('--seeds', type=_parse_num_range, help='List of random seeds', required=True)
parser_generate_ws.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
parser_generate_ws.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
parser_style_mixing_example = subparsers.add_parser('style-mixing-example', help='Generate style mixing')
parser_style_mixing_example.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
parser_style_mixing_example.add_argument('--row-seeds', type=_parse_num_range, help='Random seeds to use for image rows', required=True)
parser_style_mixing_example.add_argument('--col-seeds', type=_parse_num_range, help='Random seeds to use for image columns', required=True)
parser_style_mixing_example.add_argument('--col-styles', type=_parse_num_range, help='Style layer range (default: %(default)s)', default='0-6')
parser_style_mixing_example.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
parser_style_mixing_example.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
parser_style_mixing_gradual_change = subparsers.add_parser('style-mixing-gradual-change', help='Generate gradual change')
parser_style_mixing_gradual_change.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
parser_style_mixing_gradual_change.add_argument('--row-seeds', type=_parse_num_range, help='Random seeds to use for image rows', required=True)
parser_style_mixing_gradual_change.add_argument('--col-seeds', type=_parse_num_range, help='Random seeds to use for image columns', required=True)
parser_style_mixing_gradual_change.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
parser_style_mixing_gradual_change.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
parser_style_mixing_noise = subparsers.add_parser('style-mixing-noise', help='Pose noises')
parser_style_mixing_noise.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
parser_style_mixing_noise.add_argument('--row-seeds', type=_parse_num_range, help='Random seeds to use for image rows', required=True)
parser_style_mixing_noise.add_argument('--col-seeds', type=_parse_num_range, help='Random seeds to use for image columns', required=True)
parser_style_mixing_noise.add_argument('--col-styles', type=_parse_num_range, help='Style layer range (default: %(default)s)', default='0-6')
parser_style_mixing_noise.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
parser_style_mixing_noise.add_argument('--randomize-noise', type=bool, help='Randomize noise (default: %(default)s)', default=False)
parser_style_mixing_noise.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
parser_style_mixing_multiple = subparsers.add_parser('style-mixing-multiple', help='Generate from multiple references')
parser_style_mixing_multiple.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
parser_style_mixing_multiple.add_argument('--row-seeds', type=_parse_num_range, help='Random seeds to use for image rows', required=True)
parser_style_mixing_multiple.add_argument('--col-seeds', type=_parse_num_range, help='Random seeds to use for image columns', required=True)
parser_style_mixing_multiple.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
parser_style_mixing_multiple.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
args = parser.parse_args()
kwargs = vars(args)
subcmd = kwargs.pop('command')
if subcmd is None:
print ('Error: missing subcommand. Re-run with --help for usage.')
sys.exit(1)
sc = dnnlib.SubmitConfig()
sc.num_gpus = 1
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
sc.run_dir_root = kwargs.pop('result_dir')
sc.run_desc = subcmd
func_name_map = {
'generate-images': 'run_generator.generate_images',
'generate-ws': 'run_generator.generate_ws',
'style-mixing-example': 'run_generator.style_mixing_example',
'style-mixing-gradual-change': 'run_generator.style_mixing_gradual_change',
'style-mixing-noise': 'run_generator.style_mixing_noise',
'style-mixing-multiple': 'run_generator.style_mixing_multiple'
}
dnnlib.submit_run(sc, func_name_map[subcmd], **kwargs)
def main_v5():
submit_config = dnnlib.SubmitConfig()
# Which metrics to evaluate?
metrics = []
# metrics += [metric_defaults.fid50k_interpolate_linear_v5_z]
# metrics += [metric_defaults.prd50kv5linear_black_white]
# metrics += [metric_defaults.fid50k]
# metrics += [metric_defaults.rotation_acc_discrete10k_v5]
# metrics += [metric_defaults.classifier_consistency_v5_1000]
metrics += [metric_defaults.classifier_accuracy_model_v5]
metrics += [metric_defaults.classifier_accuracy_color_v5]
metrics += [metric_defaults.classifier_accuracy_manufacturer_v5]
metrics += [metric_defaults.classifier_accuracy_body_v5]
metrics += [metric_defaults.classifier_accuracy_rotation_v5]
metrics += [metric_defaults.classifier_accuracy_ratio_v5]
metrics += [metric_defaults.classifier_accuracy_background_v5]
# metrics += [metric_base.is50k]
# metrics += [metric_base.ppl_zfull]
# metrics += [metric_base.ppl_wend]
# metrics += [metric_base.ppl_wfull]
# metrics += [metric_base.ppl_zend]
# metrics += [metric_base.ls]
# metrics += [metric_base.dummy]
# Which networks to evaluate them on?
tasks = []
# tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ', dataset_args=EasyDict(tfrecord_dir='ffhq', shuffle_mb=0), mirror_augment=True)] # karras2019stylegan-ffhq-1024x1024.pkl
# tasks += [EasyDict(run_func_name='run_metrics.run_all_snapshots', run_id=3)]
# tasks += [EasyDict(run_func_name='run_metrics.run_all_snapshots', run_id=1)]
# baseline v5
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=48, snapshot=5000)]
# CONDITIONAL
# label mapping concat
#tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=82, snapshot=4919)]
# label dropout
#tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=327, snapshot=4999)]
# disc mod
tasks += [
EasyDict(run_func_name='run_metrics.run_snapshot',
run_id=410,
snapshot=4274)
]
# ROTATION
# interpolate v5 20_80
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=212, snapshot=4999)]
# random rotation v5
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=17, snapshot=4999)]
# int reg v5
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=323, snapshot=4999)]
# How many GPUs to use?
submit_config.num_gpus = 1
# submit_config.num_gpus = 2
# submit_config.num_gpus = 4
# submit_config.num_gpus = 8
# Execute.
submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(
config.result_dir)
submit_config.run_dir_ignore += config.run_dir_ignore
for task in tasks:
for metric in metrics:
tf.reset_default_graph()
submit_config.local.do_not_copy_source_files = True
submit_config.run_desc = '%s-%s' % (task.run_func_name,
metric.name)
if task.run_func_name.endswith('run_snapshot'):
submit_config.run_desc += '-%s-%s' % (task.run_id,
task.snapshot)
if task.run_func_name.endswith('run_all_snapshots'):
submit_config.run_desc += '-%s' % task.run_id
submit_config.run_desc += '-%dgpu' % submit_config.num_gpus
dnnlib.submit_run(submit_config, metric_args=metric, **task)
def main():
parser = argparse.ArgumentParser(
description=
'StyleGAN2 style mixer - mixes styles of row and column images.',
epilog=_examples,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument(
'--network',
default='gdrive:networks/stylegan2-ffhq-config-f.pkl',
# default='models/stylegan2-ffhq-config-f.pkl',
help='Network pickle filename',
dest='network_pkl')
parser.add_argument('--in-file',
help='Path to projected npy file',
required=True)
parser.add_argument(
'--mode',
help=
'Predifined for smile, blink, yes, no. If given rest of the options are ignored',
default='custom')
parser.add_argument('--dir-file', help='Path to direction npy file')
parser.add_argument(
'--steps',
help=
'Number of steps from start to stop. Higher the steps, smoother the interpolation',
default=30)
# TODO: instead of start and stop, accept an array of intermediate values.
parser.add_argument('--start',
help='Interpolation starts from start*direction',
type=int,
default=-15)
parser.add_argument('--stop',
help='Interpolation ends at stop*direction',
type=int,
default=15)
parser.add_argument('--reverse',
help='Should the interpolation be reversed?',
action='store_true')
parser.add_argument(
'--repeat',
help='How many times should the interpolation be repeated?',
default=2,
type=int)
parser.add_argument(
'--result-dir',
help='Root directory for run results (default: %(default)s)',
default='results',
metavar='DIR')
args = parser.parse_args()
kwargs = vars(args)
if args.dir_file:
run_desc = args.in_file.split('/')[-2] + '_' + args.dir_file.split(
'/')[-1].replace('.npy', '')
else:
run_desc = args.in_file.split('/')[-2] + '_' + args.mode
sc = dnnlib.SubmitConfig()
sc.num_gpus = 1
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
sc.run_dir_root = kwargs.pop('result_dir')
sc.run_desc = run_desc
dnnlib.submit_run(sc, 'animate.animate', **kwargs)
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma,
mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='8k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 10
sched.G_lrate_base = float(
os.environ['G_LR']) if 'G_LR' in os.environ else 0.002
sched.D_lrate_base = float(
os.environ['D_LR']) if 'D_LR' in os.environ else 0.002
sched.G_lrate_base *= float(
os.environ['G_LR_MULT']) if 'G_LR_MULT' in os.environ else 1.0
sched.D_lrate_base *= float(
os.environ['D_LR_MULT']) if 'D_LR_MULT' in os.environ else 1.0
G_opt.beta2 = float(
os.environ['G_BETA2']) if 'G_BETA2' in os.environ else 0.99
D_opt.beta2 = float(
os.environ['D_BETA2']) if 'D_BETA2' in os.environ else 0.99
print('G_lrate: %f' % sched.G_lrate_base)
print('D_lrate: %f' % sched.D_lrate_base)
print('G_beta2: %f' % G_opt.beta2)
print('D_beta2: %f' % D_opt.beta2)
sched.minibatch_size_base = int(
os.environ['BATCH_SIZE']) if 'BATCH_SIZE' in os.environ else num_gpus
sched.minibatch_gpu_base = int(
os.environ['BATCH_PER']) if 'BATCH_PER' in os.environ else 1
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
desc += '-' + dataset
resolution = int(
os.environ['RESOLUTION']) if 'RESOLUTION' in os.environ else 64
dataset_args = EasyDict(tfrecord_dir=dataset, resolution=resolution)
assert num_gpus in [
1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192
]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id != 'config-f':
G.fmap_base = D.fmap_base = 8 << 10
if 'FMAP_BASE' in os.environ:
G.fmap_base = D.fmap_base = int(os.environ['FMAP_BASE']) << 10
else:
G.fmap_base = D.fmap_base = 16 << 10 # default
print('G_fmap_base: %d' % G.fmap_base)
print('D_fmap_base: %d' % D.fmap_base)
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {
128: 0.0015,
256: 0.002,
512: 0.003,
1024: 0.003
}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def main():
parser = argparse.ArgumentParser(
description='''StyleGAN2 projector.
Run 'python %(prog)s <subcommand> --help' for subcommand help.''',
epilog=_examples,
formatter_class=argparse.RawDescriptionHelpFormatter)
subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
recon_real_images_parser = subparsers.add_parser(
'recon-real-images', help='Project real images')
recon_real_images_parser.add_argument('--network',
help='Network pickle filename',
dest='network_pkl',
required=True)
recon_real_images_parser.add_argument('--input',
help='Directory with input images',
required=True)
recon_real_images_parser.add_argument(
'--masks',
help=
'Directory with masks (inpainting only). Mask filenames should be identical to the input filenames',
required=False)
recon_real_images_parser.add_argument('--tag',
help='Tag for results directory',
default='')
recon_real_images_parser.add_argument(
'--num-snapshots',
type=int,
help='Number of snapshots (default: %(default)s)',
default=20)
recon_real_images_parser.add_argument(
'--num-gpus',
type=int,
help='Number of gpus (default: %(default)s)',
default=1)
recon_real_images_parser.add_argument(
'--result-dir',
help='Root directory for run results (default: %(default)s)',
default='results',
metavar='DIR')
recon_real_images_parser.add_argument(
'--recontype',
help=
'Type of reconstruction: "none" (normal image inversion), "super-resolution", "inpaint" (default: %(default)s)',
default='none')
recon_real_images_parser.add_argument(
'--superres-factor',
help='Super-resolution factor: 2,4,8,16,32,64 (default: %(default)s)',
type=int,
default=4)
recon_real_images_parser.add_argument(
'--num-steps',
help=
'Number of iteration steps: <1000 (fast) to 5000 (slow, but better results) (default: %(default)s)',
type=int,
default=5000)
args = parser.parse_args()
subcmd = args.command
if subcmd is None:
print('Error: missing subcommand. Re-run with --help for usage.')
sys.exit(1)
kwargs = vars(args)
sc = dnnlib.SubmitConfig()
sc.num_gpus = args.num_gpus
kwargs.pop('num_gpus')
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
sc.run_dir_root = kwargs.pop('result_dir')
print('args:', args)
command = kwargs.pop('command')
sc.run_desc = '%s-%s' % (args.tag, args.recontype)
tag = kwargs.pop('tag')
func_name_map = {'recon-real-images': 'recon.recon_real_images'}
dnnlib.submit_run(sc, func_name_map[subcmd], **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,
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_dcp=40,
hy_ncut=1,
hy_ncut_G=1,
hy_rec=20,
hy_hes=20,
hy_lin=20,
hy_mat=80,
hy_gmat=0,
hy_oth=80,
hy_det=0,
hy_1p=0,
hy_commute=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,
epsilon=1,
drange_net=[-1, 1],
recons_type='bernoulli_loss',
R_view_scale=1,
group_feat_type='concat',
normalize_alg=True,
use_alg_var=True,
use_sphere_points=False,
use_learnable_sphere_points=False,
n_sphere_points=100,
use_group_decomp=False,
mapping_after_exp=False,
snapshot_ticks=10,
subgroup_sizes_ls=None,
subspace_sizes_ls=None,
lie_alg_init_type_ls=None,
forward_eg=False,
forward_eg_prob=0.3333):
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)
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,
subgroup_sizes_ls=subgroup_sizes_ls,
subspace_sizes_ls=subspace_sizes_ls,
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,
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,
hy_ncut=hy_ncut_G,
normalize_alg=normalize_alg,
use_alg_var=use_alg_var,
subgroup_sizes_ls=subgroup_sizes_ls,
subspace_sizes_ls=subspace_sizes_ls,
lie_alg_init_type_ls=lie_alg_init_type_ls,
forward_eg=forward_eg,
forward_eg_prob=forward_eg_prob,
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.
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 == 'coma_vae': # COMA-VAE
G_loss = EasyDict(
func_name='training.loss_vae.coma_vae',
latent_type=latent_type,
hy_gamma=hy_gamma,
epsilon=epsilon,
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 == 'group_norm_vae': # GroupVAE-v3
G_loss = EasyDict(
func_name='training.loss_vae_group_v3.group_norm_vae',
latent_type=latent_type,
hy_beta=hy_beta,
hy_hes=hy_hes,
hy_commute=hy_commute,
recons_type=recons_type) # Options for generator loss.
elif model_type == 'sbs_vae': # GroupVAE-v4
G_loss = EasyDict(
func_name='training.loss_vae_group_v4.group_subspace_vae',
latent_type=latent_type,
subgroup_sizes_ls=subgroup_sizes_ls,
subspace_sizes_ls=subspace_sizes_ls,
hy_beta=hy_beta,
hy_hes=hy_hes,
hy_rec=hy_rec,
hy_commute=hy_commute,
forward_eg=forward_eg,
recons_type=recons_type) # Options for generator loss.
elif model_type == 'so_vae':
G_loss = EasyDict(
func_name='training.loss_vae_so.so_vae',
hy_1p=hy_1p,
hy_beta=hy_beta,
latent_type=latent_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,
subgroup_sizes_ls=subgroup_sizes_ls,
subspace_sizes_ls=subspace_sizes_ls,
forward_eg=forward_eg,
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 main(args):
desc = 'sgan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # 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_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict(tfrecord_dir=args.tfrecord_dirname,
resolution=args.resolution)
desc += '-ffhq' + str(args.resolution)
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='4k', layout='random') # Options for setup_snapshot_image_grid().
metrics = [metric_base.fid50k] # Options for MetricGroup.
submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.mirror_augment = True
# Number of GPUs.
submit_config.num_gpus = args.n_gpus
desc += '-{}gpu'.format(args.n_gpus)
sched.minibatch_base = 4 * args.n_gpus
sched.minibatch_dict = MINI_BATCH_DICT[args.n_gpus]
# Default options.
train.total_kimg = args.total_kimg
sched.lod_initial_resolution = 8
sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)
# kwargs
kwargs = EasyDict(train)
kwargs.update(data_root_dir=args.datadir)
kwargs.update(resume_run_id=args.resume_run_id,
resume_snapshot=args.resume_snapshot,
resume_kimg=args.resume_kimg,
resume_time=args.resume_time)
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,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.run_dir_root = args.outdir
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(model, dataset, data_dir, result_dir, config_id, num_gpus, total_kimg,
gamma, mirror_augment, metrics, max_images, resume_pkl, resume_kimg,
resume_time, pr):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.' + model +
'.G_main') # Options for generator network.
D = EasyDict(func_name='training.' + model +
'.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
train.resume_pkl = resume_pkl
train.resume_kimg = resume_kimg
train.resume_time = resume_time
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)
dataset_args.update(max_images=max_images)
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'] or\
model in ['networks_stylegan2_att', 'networks_stylegan2_satt', 'networks_stylegan2_base',
'networks_stylegan2_resample', 'networks_stylegan2_nonoise']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
if pr is not None:
if pr == 'true':
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg')
elif pr == 'false':
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
if 'dcgan' in model:
sched.G_lrate_base = sched.D_lrate_base = 0.0002
sched.minibatch_size_base = 128
sched.minibatch_gpu_base = 32 # (default)
G.func_name = 'training.dcgan.G_main'
D.func_name = 'training.dcgan.D_stylegan2'
train.run_func_name = 'training.dcgan_loop.training_loop'
G_opt = EasyDict(beta1=0.5, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.5, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss.func_name = 'training.loss.G_logistic_ns'
D_loss.func_name = 'training.loss.D_logistic'
# G_loss = EasyDict(func_name='training.loss.G_loss_dcgan')
# D_loss = EasyDict(func_name='training.loss.D_loss_dcgan') # Options for discriminator loss.
if 'add' in model:
G.noise_style = 'add'
elif 're' in model:
G.noise_style = 're'
if 'church' in dataset:
G.clip_style = 'church'
elif 'cat' in dataset:
G.clip_style = 'cat'
if 'cifar' in dataset:
G.func_name = 'training.dcgan_cifar10.G_main'
D.func_name = 'training.dcgan_cifar10.D_stylegan2'
if 'wgan' in model:
G_loss.func_name = 'training.loss.G_wgan'
D_loss.func_name = 'training.loss.D_wgan_gp'
# if model in ['networks_stylegan2_resample']:
# if 'cat' in dataset:
# G.clip_style = 'cat'
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.conditional.v5_baseline.training_loop'
) # Options for training loop.
G = EasyDict(
func_name='training.networks.conditional.label_mapping_concat.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.v5_baseline.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(
func_name='training.loss.conditional.v5_baseline.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.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = True
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_mapping_concat'
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(g_loss, g_loss_kwargs, d_loss, d_loss_kwargs, dataset_train,
dataset_eval, data_dir, result_dir, config_id, num_gpus, total_kimg,
gamma, mirror_augment, metrics, resume_pkl, resume_kimg,
resume_pkl_dir, max_images, lrate_base, img_ticks, net_ticks,
skip_images):
if g_loss_kwargs != '': g_loss_kwargs = json.loads(g_loss_kwargs)
else: g_loss_kwargs = {}
if d_loss_kwargs != '': d_loss_kwargs = json.loads(d_loss_kwargs)
else: d_loss_kwargs = {}
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(
func_name='training.loss.' + g_loss, **g_loss_kwargs
) #G_logistic_ns_gsreg') # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.' + d_loss,
**d_loss_kwargs) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='1080p',
layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = img_ticks
train.network_snapshot_ticks = net_ticks
G.scale_func = 'training.networks_stylegan2.apply_identity'
D.scale_func = None
sched.G_lrate_base = sched.D_lrate_base = lrate_base #0.002
# TODO: Changed this to 16 to match DiffAug
sched.minibatch_size_base = 16
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
sched.tick_kimg_base = 1
sched.tick_kimg_dict = {
} #{8:28, 16:24, 32:20, 64:16, 128:12, 256:8, 512:6, 1024:4}): # Resolution-specific overrides.
desc += '-' + dataset_train.split('/')[-1]
# Get dataset paths
t_path = dataset_train.split('/')
e_path = dataset_eval.split('/')
if len(t_path) > 1:
dataset_train = t_path[-1]
train.train_data_dir = os.path.join(data_dir, '/'.join(t_path[:-1]))
if len(e_path) > 1:
dataset_eval = e_path[-1]
train.eval_data_dir = os.path.join(data_dir, '/'.join(e_path[:-1]))
dataset_args = EasyDict(tfrecord_dir=dataset_train)
# Limit number of training images during train (not eval)
dataset_args['max_images'] = max_images
if max_images: desc += '-%dimg' % max_images
dataset_args['skip_images'] = skip_images
dataset_args_eval = EasyDict(tfrecord_dir=dataset_eval)
desc += '-' + dataset_eval
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
if mirror_augment: desc += '-aug'
# Infer pretrain checkpoint from target dataset
if not resume_pkl:
if any(ds in dataset_train.lower()
for ds in ['obama', 'celeba', 'rem', 'portrait']):
resume_pkl = 'ffhq-config-f.pkl'
if any(ds in dataset_train.lower()
for ds in ['gogh', 'temple', 'tower', 'medici', 'bridge']):
resume_pkl = 'church-config-f.pkl'
if any(ds in dataset_train.lower() for ds in ['bus']):
resume_pkl = 'car-config-f.pkl'
resume_pkl = os.path.join(resume_pkl_dir, resume_pkl)
train.resume_pkl = resume_pkl
train.resume_kimg = resume_kimg
train.resume_with_new_nets = True # Recreate with new parameters
# Adaptive parameters
if 'ada' in config_id:
G['train_scope'] = D[
'train_scope'] = '.*adapt' # Freeze old parameters
if 'ss' in config_id:
G['adapt_func'] = D[
'adapt_func'] = 'training.networks_stylegan2.apply_adaptive_scale_shift'
if 'sv' or 'pc' in config_id: # [:9] == 'config-sv' or config_id[:9] == 'config-pc':
G['map_svd'] = G['syn_svd'] = D['svd'] = True
# Flatten over spatial dimension
if 'flat' in config_id:
G['spatial'] = D['spatial'] = True
# Do PCA by centering before SVD
if 'pc' in config_id:
G['svd_center'] = D['svd_center'] = True
G['svd_config'] = D['svd_config'] = 'S'
if 'U' in config_id:
G['svd_config'] += 'U'
D['svd_config'] += 'U'
if 'V' in config_id:
G['svd_config'] += 'V'
D['svd_config'] += 'V'
# FreezeD
D['freeze'] = 'fd' in config_id #freeze_d
# DiffAug
if 'da' in config_id:
G_loss = EasyDict(func_name='training.loss.G_ns_diffaug')
D_loss = EasyDict(func_name='training.loss.D_ns_diffaug_r1')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(
G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss,
)
kwargs.update(dataset_args=dataset_args,
dataset_args_eval=dataset_args_eval,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(
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 main_v7():
submit_config = dnnlib.SubmitConfig()
# Which metrics to evaluate?
metrics = []
# metrics += [metric_defaults.fid50k_interpolate_linear_v7_z]
# metrics += [metric_defaults.fid50k_interpolate_spherical]
# metrics += [metric_defaults.prd10kv7spherical]
# metrics += [metric_defaults.prd50kv7spherical_black_white]
metrics += [metric_defaults.fid50k]
# metrics += [metric_defaults.rotation_acc_discrete10k_v7]
# metrics += [metric_defaults.classifier_consistency_v7]
# metrics += [metric_base.is50k]
# metrics += [metric_base.ppl_zfull]
# metrics += [metric_base.ppl_wend]
# metrics += [metric_base.ppl_wfull]
# metrics += [metric_base.ppl_zend]
# metrics += [metric_base.ls]
# metrics += [metric_base.dummy]
# Which networks to evaluate them on?
tasks = []
# tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ', dataset_args=EasyDict(tfrecord_dir='ffhq', shuffle_mb=0), mirror_augment=True)] # karras2019stylegan-ffhq-1024x1024.pkl
# tasks += [EasyDict(run_func_name='run_metrics.run_all_snapshots', run_id=3)]
tasks += [
EasyDict(run_func_name='run_metrics.run_all_snapshots', run_id=313)
]
# baseline v7
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=313, snapshot=4919)]
# interpolate cooperative v7 20_80
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=51, snapshot=4999)]
# interpolate v7 20_80
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=210, snapshot=4999)]
# random rotation v7
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=311, snapshot=4919)]
# int reg v7
# tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=2, snapshot=4354)]
# How many GPUs to use?
# submit_config.num_gpus = 1
submit_config.num_gpus = 2
# submit_config.num_gpus = 4
# submit_config.num_gpus = 8
# Execute.
submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(
config.result_dir)
submit_config.run_dir_ignore += config.run_dir_ignore
for task in tasks:
for metric in metrics:
tf.reset_default_graph()
submit_config.local.do_not_copy_source_files = True
submit_config.run_desc = '%s-%s' % (task.run_func_name,
metric.name)
if task.run_func_name.endswith('run_snapshot'):
submit_config.run_desc += '-%s-%s' % (task.run_id,
task.snapshot)
if task.run_func_name.endswith('run_all_snapshots'):
submit_config.run_desc += '-%s' % task.run_id
submit_config.run_desc += '-%dgpu' % submit_config.num_gpus
dnnlib.submit_run(submit_config, metric_args=metric, **task)
def run(**args):
args = EasyDict(args)
train = EasyDict(run_func_name = "training.training_loop.training_loop") # training loop options
sched = EasyDict() # TrainingSchedule options
vis = EasyDict() # visualize.eval() options
grid = EasyDict(size = "1080p", layout = "random") # setup_snapshot_img_grid() options
sc = dnnlib.SubmitConfig() # dnnlib.submit_run() options
# Environment configuration
tf_config = {
"rnd.np_random_seed": 1000,
"allow_soft_placement": True,
"gpu_options.per_process_gpu_memory_fraction": 1.0
}
if args.gpus != "":
num_gpus = len(args.gpus.split(","))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
# Networks configuration
cG = set_net("G", reg_interval = 4)
cD = set_net("D", reg_interval = 16)
# Dataset configuration
ratios = {
"clevr": 0.75,
"lsun-bedrooms": 0.72,
"cityscapes": 0.5,
"ffhq": 1.0
}
args.ratio = ratios.get(args.dataset, args.ratio)
dataset_args = EasyDict(tfrecord_dir = args.dataset, max_imgs = args.max_images, ratio = args.ratio,
num_threads = args.num_threads)
for arg in ["data_dir", "mirror_augment", "total_kimg"]:
cset(train, arg, args[arg])
# Training and Optimizations configuration
for arg in ["eval", "train", "recompile", "last_snapshots"]:
cset(train, arg, args[arg])
# Round to the closest multiply of minibatch size for validity
args.batch_size -= args.batch_size % args.minibatch_size
args.minibatch_std_size -= args.minibatch_std_size % args.minibatch_size
args.latent_size -= args.latent_size % args.component_num
if args.latent_size == 0:
print(bcolored("Error: latent-size is too small. Must best a multiply of component-num.", "red"))
exit()
sched_args = {
"G_lrate": "g_lr",
"D_lrate": "d_lr",
"minibatch_size": "batch_size",
"minibatch_gpu": "minibatch_size"
}
for arg, cmd_arg in sched_args.items():
cset(sched, arg, args[cmd_arg])
cset(train, "clip", args.clip)
# Logging and metrics configuration
metrics = [metric_defaults[x] for x in args.metrics]
cset(cG.args, "truncation_psi", args.truncation_psi)
for arg in ["summarize", "keep_samples"]:
cset(train, arg, args[arg])
# Visualization
args.imgs = args.images
args.ltnts = args.latents
vis_types ["imgs", "ltnts", "maps", "layer_maps", "interpolations", "noise_var", "style_mix"]:
# Set of all the set visualization types option
vis.vis_types = {arg for arg in vis_types if args[arg]}
vis_args = {
"grid": "vis_grid" ,
"num": "vis_num" ,
"rich_num": "vis_rich_num",
"section_size": "vis_section_size",
"intrp_density": "intrpolation_density",
"intrp_per_component": "intrpolation_per_component",
"alpha": "blending_alpha"
}
for arg, cmd_arg in vis_args.items():
cset(vis, arg, args[cmd_arg])
# Networks architecture
cset(cG.args, "architecture", args.g_arch)
cset(cD.args, "architecture", args.d_arch)
cset([cG.args, cD.args], "resnet_mlp", args.resnet_mlp)
cset(cG.args, "tanh", args.tanh)
# Latent sizes
if args.component_num > 1
if not (args.attention or args.merge):
print(bcolored("Error: component-num > 1 but the model is not using components.", "red"))
print(bcolored("Either add --attention for GANsformer or --merge for k-GAN).", "red"))
exit()
args.latent_size = int(args.latent_size / args.component_num)
cD.args.latent_size = cG.args.latent_size = cG.args.dlatent_size = args.latent_size
cset([cG.args, cD.args, train, vis], "component_num", args.component_num)
# Mapping network
for arg in ["layersnum", "lrmul", "dim", "shared"]:
cset(cG.args, arg, args["mapping_{}".formt(arg)])
# StyleGAN settings
for arg in ["style", "latent_stem", "fused_modconv", "local_noise"]:
cset(cG.args, arg, args[arg])
cD.args.mbstd_group_size = args.minibatch_std_size
# GANsformer
cset([cG.args, train], "attention", args.transformer)
cset(cD.args, "attention", args.d_transformer)
cset([cG.args, cD.args], "num_heads", args.num_heads)
args.norm = args.normalize
for arg in ["norm", "integration", "ltnt_gate", "img_gate", "kmeans",
"kmeans_iters", "asgn_direct", "mapping_ltnt2ltnt"]:
cset(cG.args, arg, args[arg])
for arg in ["attention_inputs", "use_pos"]:
cset([cG.args, cD.args], arg, args[arg])
# Positional encoding
for arg in ["dim", "init", "directions_num"]:
field = "pos_{}".format(arg)
cset([cG.args, cD.args], field, args[field])
# k-GAN
for arg in ["layer", "type", "channelwise"]:
field = "merge_{}".format(arg)
cset(cG.args, field, args[field])
cset([cG.args, train], "merge", args.merge)
# Attention
for arg in ["start_res", "end_res", "ltnt2ltnt", "img2img", "local_attention"]:
cset(cG.args, arg, args["g_{}".format(arg)])
cset(cD.args, arg, args["d_{}".format(arg)])
cset(cG.args, "img2ltnt", args.g_img2ltnt)
cset(cD.args, "ltnt2img", args.d_ltnt2img)
# Mixing and dropout
for arg in ["style_mixing", "component_mixing", "component_dropout", "attention_dropout"]:
cset(cG.args, arg, args[arg])
# Loss and regularization
gloss_args = {
"loss_type": "g_loss",
"reg_weight": "g_reg_weight"
"pathreg": "pathreg",
}
dloss_args = {
"loss_type": "d_loss",
"reg_type": "d_reg",
"gamma": "gamma"
}
for arg, cmd_arg in gloss_args.items():
cset(cG.loss_args, arg, args[cmd_arg])
for arg, cmd_arg in dloss_args.items():
cset(cD.loss_args, arg, args[cmd_arg])
##### Experiments management:
# Whenever we start a new experiment we store its result in a directory named 'args.expname:000'.
# When we rerun a training or evaluation command it restores the model from that directory by default.
# If we wish to restart the model training, we can set --restart and then we will store data in a new
# directory: 'args.expname:001' after the first restart, then 'args.expname:002' after the second, etc.
# Find the latest directory that matches the experiment
exp_dir = sorted(glob.glob("{}/{}:*".format(args.result_dir, args.expname)))[-1]
run_id = int(exp_dir.split(":")[-1])
# If restart, then work over a new directory
if args.restart:
run_id += 1
run_name = "{}:{0:03d}".format(args.expname, run_id)
train.printname = "{} ".format(misc.bold(args.expname))
snapshot, kimg, resume = None, 0, False
pkls = sorted(glob.glob("{}/{}/network*.pkl".format(args.result_dir, run_name)))
# Load a particular snapshot is specified
if args.pretrained_pkl:
# Soft links support
snapshot = glob.glob(args.pretrained_pkl)[0]
if os.path.islink(snapshot):
snapshot = os.readlink(snapshot)
# Extract training step from the snapshot if specified
try:
kimg = int(snapshot.split("-")[-1].split(".")[0])
except:
pass
# Find latest snapshot in the directory
elif len(pkls) > 0:
snapshot = pkls[-1]
kimg = int(snapshot.split("-")[-1].split(".")[0])
resume = True
if snapshot:
print(misc.bcolored("Resuming {}, kimg {}".format(snapshot, kimg), "white"))
train.resume_pkl = snapshot
train.resume_kimg = kimg
else:
print("Start model training from scratch.", "white")
# Run environment configuration
sc.run_dir_root = args.result_dir
sc.run_desc = args.expname
sc.run_id = run_id
sc.run_name = run_name
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(cG = cG, cD = cD)
kwargs.update(dataset_args = dataset_args, vis_args = vis, sched_args = sched, grid_args = grid, metric_arg_list = metrics, tf_config = tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.resume = resume
# If reload new options from the command line, no need to load the original configuration file
kwargs.load_config = not args.reload
dnnlib.submit_run(**kwargs)
tflib.autosummary.save_summaries(summary_log, cur_nimg)
ctx.update('%.2f' % sched.lod, cur_epoch=cur_nimg // 1000, max_epoch=total_kimg)
maintenance_time = ctx.get_last_update_interval() - tick_time
# 保存最终结果。
misc.save_pkl((g.G, g.D, g.Gs), os.path.join(submit_config.run_dir, 'network-final.pkl'))
summary_log.close()
ctx.close()
loss_values = np.array(loss_values)
np.save('./{}/loss_values.npy'.format(output_dir), loss_values)
f, ax = plt.subplots(figsize=(10, 4))
ax.plot(loss_values)
f.savefig('./{}/loss_values.png'.format(output_dir))
if __name__ == '__main__':
opt = TrainOptions().parse()
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu
# train loop
kwargs.update(opt=opt, **kwargs)
dnnlib.submit_run(**kwargs)
# loss_values = joint_train(g, graph_inputs, output_dir, opt.model_save_freq, **kwargs)
# loss_values = np.array(loss_values)
# np.save('./{}/loss_values.npy'.format(output_dir), loss_values)
# f, ax = plt.subplots(figsize=(10, 4))
# ax.plot(loss_values)
# f.savefig('./{}/loss_values.png'.format(output_dir))
def run(result_dir,
num_gpus,
total_kimg,
mirror_augment,
metrics,
resume_pkl,
G_pkl,
I_fmap_base=8,
fmap_decay=0.15,
n_samples_per=10,
module_list=None,
latent_type='uniform',
batch_size=32,
batch_per_gpu=16,
random_seed=1000,
fmap_min=16,
fmap_max=512,
dlatent_size=10,
I_nf_scale=4,
arch='resnet'):
print('module_list:', module_list)
train = EasyDict(
run_func_name='training.training_loop_infernet.training_loop_infernet'
) # Options for training loop.
module_list = _str_to_list(module_list)
I = EasyDict(func_name='training.vc_networks2.infer_modular',
dlatent_size=dlatent_size,
fmap_min=fmap_min,
fmap_max=fmap_max,
module_list=module_list,
I_nf_scale=I_nf_scale)
desc = 'inference_net'
I_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
loss = EasyDict(func_name='training.loss_inference.I_loss',
latent_type=latent_type,
dlatent_size=dlatent_size) # Options for generator loss.
sched = EasyDict() # Options for TrainingSchedule.
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
# tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
tf_config = {
'rnd.np_random_seed': random_seed
} # Options for tflib.init_tf().
train.total_kimg = total_kimg
sched.lrate = 0.002
sched.tick_kimg = 1
sched.minibatch_size = batch_size
sched.minibatch_gpu = batch_per_gpu
metrics = [metric_defaults[x] for x in metrics]
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
# Configs A-E: Shrink networks to match original StyleGAN.
I.fmap_base = 2 << I_fmap_base
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(I_args=I, I_opt_args=I_opt, loss_args=loss)
kwargs.update(sched_args=sched,
metric_arg_list=metrics,
tf_config=tf_config,
resume_pkl=resume_pkl,
G_pkl=G_pkl,
n_samples_per=n_samples_per)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
