def prepare(tasks, data_dir, max_images = None, ratio = 1.0, images_dir = None, format = None): # Options for custom dataset os.makedirs(data_dir, exist_ok = True) for task in tasks: # If task not in catalog, create custom task configuration c = catalog.get(task, EasyDict({ "local": True, "name": task, "dir": images_dir, "ratio": ratio, "process": formats_catalog.get(format) })) dirname = f"{data_dir}/{task}" os.makedirs(dirname, exist_ok = True) # try: print(misc.bold(f"Preparing the {c.name} dataset...")) if "local" not in c: fname = f"{dirname}/{c.filename}" download = not ((os.path.exists(fname) and verify_md5(fname, c.md5))) path = get_path(c.url, dirname, path = c.filename) if download: print(misc.bold(f"Downloading the data ({c.size} GB)...")) download_file(c.url, path) if path.endswith(".zip"): if not is_unzipped(path, dirname): print(misc.bold(f"Unzipping {path}...")) unzip(path, dirname) if "process" in c: imgdir = images_dir if "local" in c else (f"{dirname}/{c.dir}") c.process(dirname, imgdir, ratio = c.ratio, max_imgs = max_images) print(misc.bcolored(f"Completed preparations for {c.name}!", "blue"))
def main(): submit_config = dnnlib.SubmitConfig() # Which metrics to evaluate? metrics = [] metrics += [metric_base.fid50k] #metrics += [metric_base.ppl_zfull] #metrics += [metric_base.ppl_wfull] #metrics += [metric_base.ppl_zend] #metrics += [metric_base.ppl_wend] #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_snapshot', run_id=100, snapshot=25000)] #tasks += [EasyDict(run_func_name='run_metrics.run_all_snapshots', run_id=100)] # 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 = base_modules.ganspace.models.stylegan.stylegan_tf.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: 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(): 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, 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 = dnnlib.submission.submit.get_template_from_path( config.result_dir) kwargs.submit_config.run_dir_ignore += config.run_dir_ignore kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def main(): kwargs = EasyDict(train) kwargs.update(Encoder_args=Encoder, E_opt_args=E_opt, D_opt_args=D_opt, E_loss_args=E_loss, D_loss_args=D_loss, lr_args=lr) kwargs.update(dataset_args=Data_dir, decoder_pkl=Decoder_pkl, tf_config=tf_config) kwargs.lr_args.decay_step = train.max_iters // 4 kwargs.submit_config = copy.deepcopy(submit_config) kwargs.submit_config.num_gpus = num_gpus kwargs.submit_config.image_size = image_size kwargs.submit_config.batch_size = batch_size kwargs.submit_config.batch_size_test = batch_size_test kwargs.submit_config.z_dim = z_dim kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path( config.result_dir) kwargs.submit_config.run_dir_ignore += config.run_dir_ignore kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def easygen_train(model_path, images_path, dataset_path, start_kimg=7000, max_kimg=25000, schedule='', seed=1000): #import stylegan #from stylegan import config ##from stylegan import dnnlib #from stylegan.dnnlib import EasyDict #images_dir = '/content/raw' #max_kimg = 25000 #start_kimg = 7000 #schedule = '' #model_in = '/content/karras2019stylegan-cats-256x256.pkl' #dataset_dir = '/content/stylegan_dataset' #os.path.join(cwd, 'cache', 'stylegan_dataset') import config config.data_dir = '/content/datasets' config.results_dir = '/content/results' config.cache_dir = '/contents/cache' run_dir_ignore = [ '/contents/results', '/contents/datasets', 'contents/cache' ] import copy import dnnlib from dnnlib import EasyDict from metrics import metric_base # Prep dataset import dataset_tool print("prepping dataset...") dataset_tool.create_from_images(tfrecord_dir=dataset_path, image_dir=images_path, shuffle=False) # Set up training parameters 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() # Options for load_dataset(). sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', 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': seed} # Options for tflib.init_tf(). # Dataset desc += '-custom' dataset = EasyDict(tfrecord_dir=dataset_path) train.mirror_augment = True # Number of GPUs. desc += '-1gpu' submit_config.num_gpus = 1 sched.minibatch_base = 4 sched.minibatch_dict = { 4: 128, 8: 128, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8, 512: 4 } #{4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 16} # Default options. train.total_kimg = max_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) # schedule schedule_dict = { 4: 160, 8: 140, 16: 120, 32: 100, 64: 80, 128: 60, 256: 40, 512: 30, 1024: 20 } #{4: 2, 8:2, 16:2, 32:2, 64:2, 128:2, 256:2, 512:2, 1024:2} # Runs faster for small datasets if len(schedule) >= 5 and schedule[0] == '{' and schedule[ -1] == '}' and ':' in schedule: # is schedule a string of a dict? try: temp = eval(schedule) schedule_dict = dict(temp) # assert: it is a dict except: pass elif len(schedule) > 0: # is schedule an int? try: schedule_int = int(schedule) #assert: schedule is an int schedule_dict = {} for i in range(1, 10): schedule_dict[int(math.pow(2, i + 1))] = schedule_int except: pass print('schedule:', str(schedule_dict)) sched.tick_kimg_dict = schedule_dict # resume kimg resume_kimg = start_kimg # path to model resume_run_id = model_path # tick snapshots image_snapshot_ticks = 1 network_snapshot_ticks = 1 # Submit run 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, sched_args=sched, grid_args=grid, tf_config=tf_config) kwargs.update(resume_kimg=resume_kimg, resume_run_id=resume_run_id) kwargs.update(image_snapshot_ticks=image_snapshot_ticks, network_snapshot_ticks=network_snapshot_ticks) kwargs.submit_config = copy.deepcopy(submit_config) kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path( config.result_dir) kwargs.submit_config.run_dir_ignore += config.run_dir_ignore kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA. """Main entry point for training StyleGAN and ProGAN networks.""" import copy import dnnlib from dnnlib import EasyDict import config from metrics import metric_base #---------------------------------------------------------------------------- # Official training configs for StyleGAN, targeted mainly for FFHQ. if 1: 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() # Options for load_dataset(). sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict(
import hashlib import argparse import numpy as np from training import misc from dnnlib import EasyDict import dataset_tool catalog = { "ffhq": EasyDict({ "name": "FFHQ", # Dataset name for logging "filename": "ffhq-r08.tfrecords1of1", # Local file name "url": "http://downloads.cs.stanford.edu/nlp/data/dorarad/ffhq-r08.tfrecords1of1", # download URL "md5": "74de4f07dc7bfb07c0ad4471fdac5e67", # MD5 checksum to potentially skip download "ratio": 1, # height/width ratio "size": 13, # download size in GB "shards": 1, # Number of tfrecord shards "img_num": 70000 # Number of images }), "bedrooms": EasyDict({ "name": "LSUN-Bedrooms", # Dataset name for logging "filename": "bedroom_train_lmdb.zip", "url": "http://dl.yf.io/lsun/scenes/bedroom_train_lmdb.zip", "md5": "f2c5d904a82a6295dbdccb322b4b0a99", "dir": "bedroom_train_lmdb", "ratio": 188 / 256, "size": 43, "shards": 64,
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.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 = 'model' 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 = mirror_augment 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.v7_baseline.training_loop') G = EasyDict(func_name='training.networks.rotation.v7_baseline.G_main') D = EasyDict(func_name='training.networks.rotation.v7_baseline.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.v7_int_reg_vgg.G_logistic_ns_pathreg') D_loss = EasyDict(func_name='training.loss.rotation.v7_int_reg_vgg.D_logistic_r1') sched = EasyDict() grid = EasyDict(size='1080p', layout='random') sc = dnnlib.SubmitConfig() tf_config = {'rnd.np_random_seed': 1000} # train.resume_pkl = './../results/00422-rotation_v7_int_reg_vgg/network-snapshot-003628.pkl' # train.resume_kimg = 3628.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 = 'rotation_v7_int_reg_vgg_increase' G_loss.int_pl_weight = 20.0 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)
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA. """Main entry point for training StyleGAN networks.""" import copy import dnnlib from dnnlib import EasyDict import argparse import config from metrics import metric_base #---------------------------------------------------------------------------- # Official training configs for StyleGAN, targeted mainly for FFHQ. if 1: desc = 'sgan' train = EasyDict() # Description string included in result subdir name. 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() # Options for load_dataset(). sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict(
import copy import config import dnnlib from dnnlib import EasyDict from metrics import metric_base # ---------------------------------------------------------------------------- # Official training configs for StyleGAN, targeted mainly for FFHQ. if 1: # Description string included in result subdir name. desc = 'sgan' # Options for training loop. train = EasyDict(run_func_name='training.training_loop.training_loop') # Options for generator network. G = EasyDict(func_name='training.networks_stylegan.G_style') # Options for discriminator network. D = EasyDict(func_name='training.networks_stylegan.D_basic') # Options for generator optimizer. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator loss. G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating') # Options for discriminator loss. D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp', r1_gamma=10.0) # Options for load_dataset(). dataset = EasyDict()
# Default metric definitions from dnnlib import EasyDict metric_defaults = EasyDict([(args.name, args) for args in [ EasyDict(name = "fid", func_name = "metrics.frechet_inception_distance.FID", minibatch_per_gpu = 8), EasyDict(name = "is", func_name = "metrics.inception_score.IS", num_splits = 10, minibatch_per_gpu = 8), EasyDict(name = "pr", func_name = "metrics.precision_recall.PR", nhood_size = 3, minibatch_per_gpu = 8, row_batch_size = 5000, col_batch_size = 5000), # EasyDict(name = "ppl_zfull", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "z", sampling = "full", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")), # EasyDict(name = "ppl_wfull", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "w", sampling = "full", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")), # EasyDict(name = "ppl_zend", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "z", sampling = "end", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")), # EasyDict(name = "ppl_wend", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "w", sampling = "end", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")), # EasyDict(name = "ppl2_wend", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "w", sampling = "end", crop = False, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")), # EasyDict(name = "ls", func_name = "metrics.linear_separability.LS", num_samples = 200000, num_keep = 100000, attrib_indices = range(40), minibatch_per_gpu = 4), ]])
import copy import os import dnnlib from dnnlib import EasyDict import config from metrics import metric_base # ---------------------------------------------------------------------------- # Official training configs for StyleGAN, targeted mainly for FFHQ. # turn off black formatting for this file: # fmt: off desc = 'msg-stylegan' # 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() # Options for load_dataset(). sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict(size='4k',
import os import dnnlib from dnnlib import EasyDict import config import copy desc = 'stylegan-encoder' + config.DESC train = EasyDict(run_func_name='training.training_loop_encoder.training_loop') Encoder = EasyDict(func_name='training.networks_encoder.Encoder') E_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8) D_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8) E_loss = EasyDict(func_name='training.loss_encoder.E_loss', feature_scale=0.00005, D_scale=0.05, perceptual_img_size=256) D_loss = EasyDict(func_name='training.loss_encoder.D_logistic_simplegp', r1_gamma=10.0) lr = EasyDict(learning_rate=0.0001, decay_step=30000, decay_rate=0.8, stair=False) Data_dir = EasyDict(data_train=config.DATA_DIR, data_test=config.TEST_DATA_DIR) Decoder_pkl = EasyDict(decoder_pkl=config.PICKLE_DIR) Inversion_pkl = EasyDict(inversion_pkl=config.INVERSION_PICKLE_DIR) tf_config = {'rnd.np_random_seed': 1000} submit_config = dnnlib.SubmitConfig() num_gpus = 1 desc += '-1gpu' # num_gpus = 2; desc += '-2gpu' # num_gpus = 4; desc += '-4gpu'
import os import dnnlib from dnnlib import EasyDict import config import copy #os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" desc = 'LIA' train = EasyDict(run_func_name='training.training_loop_encoder.training_loop') Encoder = EasyDict(func_name='training.networks_encoder.Encoder') E_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8) D_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8) E_loss = EasyDict(func_name='training.loss_encoder.E_perceptual_loss', feature_scale=0.00005, D_scale=0.1) D_loss = EasyDict(func_name='training.loss_encoder.D_logistic_simplegp', r1_gamma=10.0) lr = EasyDict(learning_rate=0.0001, decay_step=50000, decay_rate=0.8, stair=False) Data_dir = EasyDict(data_train='../data/ffhq_128_train_6w5.tfrecords', data_test='../data/ffhq_128_test_5k.tfrecords') Decoder_pkl = EasyDict(decoder_pkl='./network-snapshot-023315.pkl' ) #the first stage training results tf_config = {'rnd.np_random_seed': 1000} submit_config = dnnlib.SubmitConfig() # num_gpus = 1; desc += '-1gpu' # num_gpus = 2; desc += '-2gpu'
"""Main entry point for training StyleGAN and ProGAN networks.""" import copy import dnnlib from dnnlib import EasyDict import config from metrics import metric_base #---------------------------------------------------------------------------- # Official training configs for StyleGAN, targeted mainly for FFHQ. if 1: 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() # Options for load_dataset(). sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict(size='4k', layout='random') # Options for setup_snapshot_image_grid(). metrics = [metric_base.fid50k] # Options for MetricGroup. submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). # Dataset. desc += '-images_512x512'; dataset = EasyDict(tfrecord_dir='images_512x512'); train.mirror_augment = True
def main(): 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, 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 = dnnlib.submission.submit.get_template_from_path(config.result_dir) kwargs.submit_config.run_dir_ignore += config.run_dir_ignore kwargs.submit_config.run_desc = desc kwargs.save_weight_histograms = True kwargs.network_snapshot_ticks = 5 kwargs.update(resume_run_id=-1) dnnlib.submit_run(**kwargs)
def create_model(config_id='config-f', gamma=None, height=512, width=512, cond=None, label_size=0): train = EasyDict(run_func_name='training.diagnostic.create_initial_pkl' ) # 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. D_loss = EasyDict(func_name='training.loss.D_logistic_r1' ) # Options for discriminator 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(). sched.minibatch_size_base = 192 sched.minibatch_gpu_base = 3 D_loss.gamma = 10 desc = 'stylegan2' dataset_args = EasyDict() # (tfrecord_dir=dataset) if cond: desc += '-cond' dataset_args.max_label_size = 'full' # conditioned on full label 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 G.update(resolution_h=height) G.update(resolution_w=width) D.update(resolution_h=height) D.update(resolution_w=width) sc.submit_target = dnnlib.SubmitTarget.DIAGNOSTIC sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) # [EDITED] kwargs.update(G_args=G, D_args=D, tf_config=tf_config, config_id=config_id, resolution_h=height, resolution_w=width, label_size=label_size) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_desc = desc dnnlib.submit_diagnostic(**kwargs) return f'network-initial-config-f-{height}x{width}-{label_size}.pkl'
def main(): #------------------------------------------------------------------------ # Modified by Deng et al. args = parse_args() if args is None: exit() weight_args = EasyDict() weight_args.update(w_id=args.w_id, w_lm=args.w_lm, w_gamma=args.w_gamma, w_skin=args.w_skin, w_exp_warp=args.w_exp_warp, w_gamma_change=args.w_gamma_change) train.update(run_func_name='training.training_loop.training_loop') kwargs = EasyDict(train) # stage 1: training with only imitative losses with 15000k images. if args.stage == 1: train_stage = EasyDict( func_name='training.training_utils.training_stage1') kwargs.update(total_kimg=15000) # stage 2: training with imitative losses and contrastive losses. else: train_stage = EasyDict( func_name='training.training_utils.training_stage2') kwargs.update(resume_run_id=args.run_id, resume_snapshot=args.snapshot, resume_kimg=args.kimg) kwargs.update(total_kimg=25000) weight_args.update(w_lm=100) kwargs.update(train_stage_args=train_stage) kwargs.update(weight_args=weight_args, noise_dim=args.noise_dim) #------------------------------------------------------------------------ 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 = dnnlib.submission.submit.get_template_from_path( config.result_dir) kwargs.submit_config.run_dir_ignore += config.run_dir_ignore kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def main(): parser = argparse.ArgumentParser( description='Train ProgressiveGAN w/ Hessian Penalty') # Model/Dataset Parameters: parser.add_argument('--num_gpus', type=int, required=True, help='Number of GPUs to train with') parser.add_argument( '--dataset', type=str, default='edges_and_shoes', help='Name of TFRecords directory in datasets/ folder to train on.') parser.add_argument( '--resolution', type=int, default=128, help='Maximum resolution (after progressive growth) to train at.') parser.add_argument('--nz', type=int, default=12, help='Number of components in G\'s latent space.') parser.add_argument('--total_kimg', type=int, default=30000, help='How long to train for') parser.add_argument( '--resume_exp', type=int, default=None, help= 'If specified, resumes training from a snapshot in the results directory with job number resume_exp' ) parser.add_argument( '--resume_snapshot', default='latest', help= 'If using resume_exp, you can override this default to resume training from a specific checkpoint' ) parser.add_argument('--seed', type=int, default=1000, help='NumPy Random Seed') # Hessian Penalty Parameters: parser.add_argument( '--hp_lambda', type=float, default=0.1, help='Loss weighting of the Hessian Penalty regularization term. ' 'When fine-tuning with the Hessian Penalty, this value ' 'is the maximum weighting once the loss term is fully ramped-up.' 'Set to 0.0 to disable the Hessian Penalty.') parser.add_argument( '--epsilon', type=float, default=0.1, help='The granularity of the finite differences approximation. ' 'When changing this value from 0.1, you will likely need to change ' 'hp_lambda as well to get optimal results.') parser.add_argument( '--num_rademacher_samples', type=int, default=2, help='The number of Rademacher vectors to be sampled per-batch element ' 'when estimating the Hessian Penalty. Must be >=2. Setting this ' 'parameter larger can result in GPU out-of-memory for smaller GPUs!') parser.add_argument( '--layers_to_reg', nargs='*', default=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], help= 'Which activations to directly apply the Hessian Penalty to? For 128x128 ' 'resolution, you can choose any layers 1-13. ' '(1=first fully-connected layer, 2-12=convolutional layers, 13=RGB activations)' ) parser.add_argument( '--warmup_kimg', type=int, default=10000, help='Over how many kimg (=thousands of real images fed to D) should ' 'the Hessian Penalty weight be linearly ramped-up? Setting this ' 'value too low can sometimes result in training instability when fine-tuning.' 'The warmup will start at hp_start_kimg.') parser.add_argument( '--hp_start_kimg', type=int, default=0, help='To train from scratch with the Hessian Penalty, set this to 0. ' 'If >0, the Hessian Penalty will be phased-in starting at the ' 'specified iteration. Alternatively, you can load ' 'a checkpoint and simply set this to 0 to fine-tune a pretrained model.' 'Note that we haven\'t tested setting this > 0.') # InfoGAN Parameters: parser.add_argument( '--infogan_lambda', type=float, default=0.0, help= 'Loss weighting for InfoGAN loss on latent vectors. Set to 0.0 to disable InfoGAN loss.' ) parser.add_argument( '--infogan_nz', type=int, default=0, help= 'Number of Z components to regularize with InfoGAN\'s mutual information loss. ' 'For a (batch_size, nz) Z vector, you can access the regularized components with ' 'Z[:, :infogan_nz].') # Visualization and Metric Parameters: parser.add_argument( '--compute_interp_ticks', type=int, default=5, help= 'Training ticks between computing disentanglement visualizations (interpolations).' ) parser.add_argument( '--compute_metrics_ticks', type=int, default=5, help= 'Training ticks between computing FID (and optionally PPL; see next parameter).' ) parser.add_argument('--metrics', nargs='*', default=['FID', 'PPL'], help='Which metrics to compute during training.') parser.add_argument( '--dashboard_api', type=str, choices=['tensorboard', 'wandb'], default='tensorboard', help= 'Which training dashboard software to use for logging visualizations/metrics/losses/etc; ' 'either TensorBoard (default) or WandB. If you choose to use WandB, you will first need ' 'an account on wandb.com.') parser.add_argument('--wandb_entity', type=str, default=None, help='If using WandB, the entity for logging' '(e.g., your username)') opt = parser.parse_args() opt = EasyDict(vars(opt)) verify_opt(opt) run(opt)
def run(opt): """Sets-up all of the parameters necessary to start a ProgressiveGAN training job.""" desc = build_job_name( opt) # 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_progan.G_paper' ) # Options for generator network. D = EasyDict(func_name='training.networks_progan.D_paper' ) # 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_wgan') # Options for generator loss. D_loss = EasyDict( func_name='training.loss.D_wgan_gp') # Options for discriminator loss. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', layout='random') # Options for setup_snapshot_image_grid(). submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = { 'rnd.np_random_seed': opt.seed } # Options for tflib.init_tf(). metrics = [] # Metrics to run during training. if 'FID' in opt.metrics: metrics.append(metric_base.fid50k) if 'PPL' in opt.metrics: metrics.append(metric_base.ppl_zend_v2) train.network_metric_ticks = opt.compute_metrics_ticks train.interp_snapshot_ticks = opt.compute_interp_ticks find_dataset(opt.dataset) # Optionally resume from checkpoint: if opt.resume_exp is not None: results_dir = os.path.join(os.getcwd(), config.result_dir) _resume_pkl = find_pkl(results_dir, opt.resume_exp, opt.resume_snapshot) train.resume_run_id = opt.resume_exp train.resume_snapshot = _resume_pkl train.resume_kimg = int(_resume_pkl.split('.pkl')[0][-6:]) if f'hessian_penalty_{opt.dataset}' not in _resume_pkl and opt.hp_lambda > 0: print( 'When fine-tuning a job that was originally trained without the Hessian Penalty, ' 'hp_start_kimg is relative to the kimg of the checkpoint being resumed from. ' 'Hessian Penalty will be phased-in starting at absolute ' f'kimg={opt.hp_start_kimg + train.resume_kimg}.') opt.hp_start_kimg += train.resume_kimg # Set up dataset hyper-parameters: dataset = EasyDict(tfrecord_dir=os.path.join(os.getcwd(), config.data_dir, opt.dataset), resolution=opt.resolution) train.mirror_augment = False # Set up network hyper-parameters: G.latent_size = opt.nz D.infogan_nz = opt.infogan_nz G.infogan_lambda = opt.infogan_lambda D.infogan_lambda = opt.infogan_lambda # When computing the multi-layer Hessian Penalty, we retrieve intermediate activations by accessing the # corresponding tensor's name. Below are the names of various activations in G that we can retrieve: activation_type = 'norm' progan_generator_layer_index_to_name = { 1: f'4x4/Dense/Post_{activation_type}', 2: f'4x4/Conv/Post_{activation_type}', 3: f'8x8/Conv0_up/Post_{activation_type}', 4: f'8x8/Conv1/Post_{activation_type}', 5: f'16x16/Conv0_up/Post_{activation_type}', 6: f'16x16/Conv1/Post_{activation_type}', 7: f'32x32/Conv0_up/Post_{activation_type}', 8: f'32x32/Conv1/Post_{activation_type}', 9: f'64x64/Conv0_up/Post_{activation_type}', 10: f'64x64/Conv1/Post_{activation_type}', 11: f'128x128/Conv0_up/Post_{activation_type}', 12: f'128x128/Conv1/Post_{activation_type}', 13: 'images_out' # final full-resolution RGB activation } # Convert from layer indices to layer names (which we'll need to compute the Hessian Penalty): layers_to_reg = [ progan_generator_layer_index_to_name[layer_ix] for layer_ix in sorted(opt.layers_to_reg) ] # Store the Hessian Penalty parameters in their own dictionary: HP = EasyDict(hp_lambda=opt.hp_lambda, epsilon=opt.epsilon, num_rademacher_samples=opt.num_rademacher_samples, layers_to_reg=layers_to_reg, warmup_nimg=opt.warmup_kimg * 1000, hp_start_nimg=opt.hp_start_kimg * 1000) # How long to train for (as measured by thousands of real images processed, not gradient steps): train.total_kimg = opt.total_kimg # We ran the original experiments using 4 GPUs per job. If using a different number, # we try to scale batch sizes up or down accordingly in the for-loop below. Note that # using other batch sizes is somewhat untested, though! submit_config.num_gpus = opt.num_gpus sched.minibatch_base = 32 sched.minibatch_dict = { 4: 2048, 8: 1024, 16: 512, 32: 256, 64: 128, 128: 96, 256: 32, 512: 16 } for res, batch_size in sched.minibatch_dict.items(): sched.minibatch_dict[res] = int(batch_size * opt.num_gpus / 4) # Set-up WandB if optionally using it instead of TensorBoard: if opt.dashboard_api == 'wandb': init_wandb(opt=opt, name=desc, group=opt.dataset, entity=opt.wandb_entity) # Start the training job: kwargs = EasyDict(train) kwargs.update(HP_args=HP, 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 = dnnlib.submission.submit.get_template_from_path( config.result_dir) kwargs.submit_config.run_dir_ignore += config.run_dir_ignore 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 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)
metric_defaults = EasyDict([(args.name, args) for args in [ EasyDict(name='fid200-rt-shoes', func_name='metrics.frechet_inception_distance.FID', num_images=200, minibatch_per_gpu=1, ref_train=True, ref_samples=49825), EasyDict(name='fid200-rt-handbags', func_name='metrics.frechet_inception_distance.FID', num_images=200, minibatch_per_gpu=1, ref_train=True, ref_samples=138567), EasyDict(name='fid5k', func_name='metrics.frechet_inception_distance.FID', num_images=5000, minibatch_per_gpu=8), EasyDict(name='fid10k', func_name='metrics.frechet_inception_distance.FID', num_images=10000, minibatch_per_gpu=8), EasyDict(name='fid10k-b1', func_name='metrics.frechet_inception_distance.FID', num_images=10000, minibatch_per_gpu=1), EasyDict(name='fid10k-h0', func_name='metrics.frechet_inception_distance.FID', num_images=10000, minibatch_per_gpu=8, hole_range=[0.0, 0.2]), EasyDict(name='fid10k-h1', func_name='metrics.frechet_inception_distance.FID', num_images=10000, minibatch_per_gpu=8, hole_range=[0.2, 0.4]), EasyDict(name='fid10k-h2', func_name='metrics.frechet_inception_distance.FID', num_images=10000, minibatch_per_gpu=8, hole_range=[0.4, 0.6]), EasyDict(name='fid10k-h3', func_name='metrics.frechet_inception_distance.FID', num_images=10000, minibatch_per_gpu=8, hole_range=[0.6, 0.8]), EasyDict(name='fid10k-h4', func_name='metrics.frechet_inception_distance.FID', num_images=10000, minibatch_per_gpu=8, hole_range=[0.8, 1.0]), EasyDict(name='fid36k5', func_name='metrics.frechet_inception_distance.FID', num_images=36500, minibatch_per_gpu=8), EasyDict(name='fid36k5-h0', func_name='metrics.frechet_inception_distance.FID', num_images=36500, minibatch_per_gpu=8, hole_range=[0.0, 0.2]), EasyDict(name='fid36k5-h1', func_name='metrics.frechet_inception_distance.FID', num_images=36500, minibatch_per_gpu=8, hole_range=[0.2, 0.4]), EasyDict(name='fid36k5-h2', func_name='metrics.frechet_inception_distance.FID', num_images=36500, minibatch_per_gpu=8, hole_range=[0.4, 0.6]), EasyDict(name='fid36k5-h3', func_name='metrics.frechet_inception_distance.FID', num_images=36500, minibatch_per_gpu=8, hole_range=[0.6, 0.8]), EasyDict(name='fid36k5-h4', func_name='metrics.frechet_inception_distance.FID', num_images=36500, minibatch_per_gpu=8, hole_range=[0.8, 1.0]), EasyDict(name='fid50k', func_name='metrics.frechet_inception_distance.FID', num_images=50000, minibatch_per_gpu=8), EasyDict(name='ids5k', func_name='metrics.inception_discriminator_score.IDS', num_images=5000, minibatch_per_gpu=8), EasyDict(name='ids10k', func_name='metrics.inception_discriminative_score.IDS', num_images=10000, minibatch_per_gpu=8), EasyDict(name='ids10k-b1', func_name='metrics.inception_discriminative_score.IDS', num_images=10000, minibatch_per_gpu=1), EasyDict(name='ids10k-h0', func_name='metrics.inception_discriminative_score.IDS', num_images=10000, minibatch_per_gpu=8, hole_range=[0.0, 0.2]), EasyDict(name='ids10k-h1', func_name='metrics.inception_discriminative_score.IDS', num_images=10000, minibatch_per_gpu=8, hole_range=[0.2, 0.4]), EasyDict(name='ids10k-h2', func_name='metrics.inception_discriminative_score.IDS', num_images=10000, minibatch_per_gpu=8, hole_range=[0.4, 0.6]), EasyDict(name='ids10k-h3', func_name='metrics.inception_discriminative_score.IDS', num_images=10000, minibatch_per_gpu=8, hole_range=[0.6, 0.8]), EasyDict(name='ids10k-h4', func_name='metrics.inception_discriminative_score.IDS', num_images=10000, minibatch_per_gpu=8, hole_range=[0.8, 1.0]), EasyDict(name='ids36k5', func_name='metrics.inception_discriminative_score.IDS', num_images=36500, minibatch_per_gpu=8), EasyDict(name='ids36k5-h0', func_name='metrics.inception_discriminative_score.IDS', num_images=36500, minibatch_per_gpu=8, hole_range=[0.0, 0.2]), EasyDict(name='ids36k5-h1', func_name='metrics.inception_discriminative_score.IDS', num_images=36500, minibatch_per_gpu=8, hole_range=[0.2, 0.4]), EasyDict(name='ids36k5-h2', func_name='metrics.inception_discriminative_score.IDS', num_images=36500, minibatch_per_gpu=8, hole_range=[0.4, 0.6]), EasyDict(name='ids36k5-h3', func_name='metrics.inception_discriminative_score.IDS', num_images=36500, minibatch_per_gpu=8, hole_range=[0.6, 0.8]), EasyDict(name='ids36k5-h4', func_name='metrics.inception_discriminative_score.IDS', num_images=36500, minibatch_per_gpu=8, hole_range=[0.8, 1.0]), EasyDict(name='ids50k', func_name='metrics.inception_discriminative_score.IDS', num_images=50000, minibatch_per_gpu=8), EasyDict( name='lpips2k', func_name='metrics.learned_perceptual_image_patch_similarity.LPIPS', num_pairs=2000, minibatch_per_gpu=8), ]])
sess = tf.get_default_session() sess.list_devices() cores = tflex.get_cores() tflex.set_override_cores(cores) #synthesis_func = 'G_synthesis_stylegan2' #kwargs = {'resolution': 512} #synthesis = tflib.Network('G_synthesis', func_name=globals()[synthesis_func], **kwargs) #sess.reset(os.environ['TPU_NAME']) # don't do this, this breaks the session train = EasyDict(run_func_name='training.training_loop.training_loop') # Options for training loop. G_args = EasyDict(func_name='training.networks_stylegan2.G_main') # Options for generator network. D_args = 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} label_dtype = np.int64 sched.minibatch_gpu = 1 if 'G' not in globals(): with tflex.device('/gpu:0'):
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, 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)
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA. """Main entry point for training StyleGAN and ProGAN networks.""" import copy import dnnlib from dnnlib import EasyDict import config from metrics import metric_base # ---------------------------------------------------------------------------- # Official training configs for StyleGAN, targeted mainly for FFHQ. if 1: 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() # Options for load_dataset(). sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict(
def G_synthesis_vc_modular( dlatents_withl_in, # Input: Disentangled latents (W) [minibatch, label_size+dlatent_size]. dlatent_size=7, # Disentangled latent (W) dimensionality. Including discrete info, rotation, scaling, xy shearing, and xy translation. label_size=0, # Label dimensionality, 0 if no labels. module_list=None, # A list containing module names, which represent semantic latents (exclude labels). num_channels=1, # Number of output color channels. resolution=64, # Output resolution. fmap_base=16 << 10, # Overall multiplier for the number of feature maps. fmap_decay=1.0, # log2 feature map reduction when doubling the resolution. fmap_min=1, # Minimum number of feature maps in any layer. fmap_max=512, # Maximum number of feature maps in any layer. nonlinearity='lrelu', # Activation function: 'relu', 'lrelu', etc. dtype='float32', # Data type to use for activations and outputs. resample_kernel=[ 1, 3, 3, 1 ], # Low-pass filter to apply when resampling activations. None = no filtering. fused_modconv=True, # Implement modulated_conv2d_layer() as a single fused op? use_noise=False, # If noise is used in this dataset. randomize_noise=True, # True = randomize noise inputs every time (non-deterministic), False = read noise inputs from variables. single_const=True, # If only use a single constant feature at the begining. where_feat_map=15, # For F_loss, which layer of feat map to use. **_kwargs): # Ignore unrecognized keyword args. ''' Modularized variation-consistent network. ''' resolution_log2 = int(np.log2(resolution)) # == 6 for resolution 64 assert resolution == 2**resolution_log2 and resolution >= 4 def nf(stage): return np.clip(int(fmap_base / (2.0**(stage * fmap_decay))), fmap_min, fmap_max) num_layers = resolution_log2 * 2 - 2 # == 10 for resolution 64 act = nonlinearity images_out = None # Note that module_list may include modules not containing latents, # e.g. Conv layers (size in this case means number of conv layers). key_ls, size_ls, count_dlatent_size, n_content = split_module_names( module_list) print('In key_ls:', key_ls) print('In size_ls:', size_ls) print('In count_dlatent_size:', count_dlatent_size) if label_size > 0: key_ls.insert(0, 'Label') size_ls.insert(0, label_size) n_content += label_size # module_dict = collections.OrderedDict(zip(key_ls, size_ls)) # Primary inputs. assert dlatent_size == count_dlatent_size dlatents_withl_in.set_shape([None, label_size + count_dlatent_size]) dlatents_withl_in = tf.cast(dlatents_withl_in, dtype) # Early layers consists of 4x4 constant layer. y = None if single_const: with tf.variable_scope('4x4'): with tf.variable_scope('Const'): x = tf.get_variable( 'const', shape=[1, 8, 4, 4], initializer=tf.initializers.random_normal()) x = tf.tile(tf.cast(x, dtype), [tf.shape(dlatents_withl_in)[0], 1, 1, 1]) else: with tf.variable_scope('4x4'): with tf.variable_scope('Const'): x = tf.get_variable( 'const', shape=[n_content, 8, 4, 4], initializer=tf.initializers.random_normal()) subkwargs = EasyDict() subkwargs.update(dlatents_withl_in=dlatents_withl_in, n_content=n_content, act=act, dtype=dtype, resample_kernel=resample_kernel, fused_modconv=fused_modconv, use_noise=use_noise, randomize_noise=randomize_noise) # Build modules by module_dict. start_idx = 0 # print('module_dict:', module_dict) # for scope_idx, k in enumerate(module_dict): for scope_idx, k in enumerate(key_ls): if scope_idx == where_feat_map: feat_map = x if (k.startswith('Label')) or (k.startswith('D_global')): # e.g. {'Label': 3}, {'D_global': 3} x = build_D_layers(x, name=k, n_latents=size_ls[scope_idx], start_idx=start_idx, scope_idx=scope_idx, single_const=single_const, fmaps=nf(scope_idx), **subkwargs) start_idx += size_ls[scope_idx] elif k.startswith('C_global'): # e.g. {'C_global': 2} x = build_C_global_layers(x, name=k, n_latents=size_ls[scope_idx], start_idx=start_idx, scope_idx=scope_idx, **subkwargs) start_idx += size_ls[scope_idx] elif k.startswith('C_nocond_global'): # e.g. {'C_nocond_global': 2} x = build_C_global_nocond_layers(x, name=k, n_latents=size_ls[scope_idx], start_idx=start_idx, scope_idx=scope_idx, fmaps=nf(scope_idx), **subkwargs) start_idx += size_ls[scope_idx] elif k.startswith('C_local_heat'): # e.g. {'C_local_heat': 4} x = build_local_heat_layers(x, name=k, n_latents=size_ls[scope_idx], start_idx=start_idx, scope_idx=scope_idx, fmaps=nf(scope_idx), **subkwargs) start_idx += size_ls[scope_idx] elif k.startswith('C_local_hfeat'): # e.g. {'C_local_hfeat_size': 4} x = build_local_hfeat_layers(x, name=k, n_latents=size_ls[scope_idx], start_idx=start_idx, scope_idx=scope_idx, fmaps=nf(scope_idx), **subkwargs) start_idx += size_ls[scope_idx] elif k.startswith('Noise'): # e.g. {'Noise': 1} x = build_noise_layer(x, name=k, n_layers=size_ls[scope_idx], scope_idx=scope_idx, fmaps=nf(scope_idx), **subkwargs) elif k.startswith('Conv'): # e.g. {'Conv-up': 2}, {'Conv-id': 1} x = build_conv_layer(x, name=k, n_layers=size_ls[scope_idx], scope_idx=scope_idx, fmaps=nf(scope_idx), **subkwargs) else: raise ValueError('Unsupported module type: ' + k) y = torgb(x, y, num_channels=num_channels) images_out = y assert images_out.dtype == tf.as_dtype(dtype) return tf.identity(images_out, name='images_out'), tf.identity(feat_map, name='feat_map')