def run_auto(dataset, data_dir, result_dir, config_id, num_gpus, resolution,
total_kimg, gamma, mirror_augment, metrics, train_auto):
train = EasyDict(run_func_name='training.training_loop.training_auto_loop'
) # Options for training loop.
Enc = EasyDict(func_name='training.networks_stylegan2.Encoder'
) # Options for encoder network.
Dec = EasyDict(func_name='training.networks_stylegan2.Decoder'
) # Options for decoder network.
opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for autoencoder optimizer.
loss = EasyDict(
func_name='training.loss.auto_l1') # Options for autoencoder loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='1080p',
layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.data_dir = data_dir
train.total_kimg = total_kimg
train.image_snapshot_ticks = 10
train.network_snapshot_ticks = 125
sched.lrate = 0.003
sched.minibatch_size = 64
sched.minibatch_gpu = 64
desc = 'stylegan2-hrae'
desc += '-' + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
dataset_args.resolution = resolution
dataset_args.num_threads = 4
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
desc += '-auto'
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(Enc_args=Enc, Dec_args=Dec, opt_args=opt, loss_args=loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma,
mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.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().
try:
pkl, kimg = misc.locate_latest_pkl(result_dir)
train.resume_pkl = pkl
train.resume_kimg = kimg
except:
print('Couldn\'t find valid snapshot, starting over')
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = 1
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, 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.image_snapshot_ticks = 1
train.network_snapshot_ticks = 4
sched.minibatch_gpu_base = 8
# preprocess dataset into tfrecords if necessary
dataset = dataset_tool.create_dataset(dataset, resolution)
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
metrics = [metric_defaults[x] for x in metrics]
metric_args = EasyDict(cache_dir=dataset, num_repeats=num_repeats)
desc = 'DiffAugment-stylegan2' if DiffAugment else 'stylegan2'
dataset_args = EasyDict(tfrecord_dir=dataset,
resolution=resolution,
from_tfrecords=True)
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(dataset, train_dir, config, d_aug, diffaug_policy, cond, ops, jpg_data, mirror, mirror_v, \
lod_step_kimg, batch_size, resume, resume_kimg, finetune, num_gpus, ema_kimg, gamma, freezeD):
# dataset (tfrecords) - preprocess or get
tfr_files = file_list(os.path.dirname(dataset), 'tfr')
tfr_files = [f for f in tfr_files if basename(dataset) in f]
if len(tfr_files) == 0:
tfr_file, total_samples = create_from_images(dataset, jpg=jpg_data)
else:
tfr_file = tfr_files[0]
dataset_args = EasyDict(tfrecord=tfr_file, jpg_data=jpg_data)
desc = basename(tfr_file).split('-')[0]
# training functions
if d_aug: # https://github.com/mit-han-lab/data-efficient-gans
train = EasyDict(
run_func_name='training.training_loop_diffaug.training_loop'
) # Options for training loop (Diff Augment method)
loss_args = EasyDict(
func_name='training.loss_diffaug.ns_DiffAugment_r1',
policy=diffaug_policy) # Options for loss (Diff Augment method)
else: # original nvidia
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop (original from NVidia)
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
# network functions
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='1080p',
layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
G.impl = D.impl = ops
# resolutions
data_res = basename(tfr_file).split('-')[-1].split(
'x') # get resolution from dataset filename
data_res = list(reversed([int(x)
for x in data_res])) # convert to int list
init_res, resolution, res_log2 = calc_init_res(data_res)
if init_res != [4, 4]:
print(' custom init resolution', init_res)
G.init_res = D.init_res = list(init_res)
train.setname = desc + config
desc = '%s-%d-%s' % (desc, resolution, config)
# training schedule
sched.lod_training_kimg = lod_step_kimg
sched.lod_transition_kimg = lod_step_kimg
train.total_kimg = lod_step_kimg * res_log2 * 2 # a la ProGAN
if finetune:
train.total_kimg = 15000 # should start from ~10k kimg
train.image_snapshot_ticks = 1
train.network_snapshot_ticks = 5
train.mirror_augment = mirror
train.mirror_augment_v = mirror_v
# learning rate
if config == 'e':
if finetune: # uptrain 1024
sched.G_lrate_base = 0.001
else: # train 1024
sched.G_lrate_base = 0.001
sched.G_lrate_dict = {0: 0.001, 1: 0.0007, 2: 0.0005, 3: 0.0003}
sched.lrate_step = 1500 # period for stepping to next lrate, in kimg
if config == 'f':
# sched.G_lrate_base = 0.0003
sched.G_lrate_base = 0.001
sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway
sched.minibatch_gpu_base = batch_size
sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base
sc.num_gpus = num_gpus
if config == 'e':
G.fmap_base = D.fmap_base = 8 << 10
if d_aug: loss_args.gamma = 100 if gamma is None else gamma
else: D_loss.gamma = 100 if gamma is None else gamma
elif config == 'f':
G.fmap_base = D.fmap_base = 16 << 10
else:
print(' Only configs E and F are implemented')
exit()
if cond:
desc += '-cond'
dataset_args.max_label_size = 'full' # conditioned on full label
if freezeD:
D.freezeD = True
train.resume_with_new_nets = True
if d_aug:
desc += '-daug'
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
tf_config=tf_config)
kwargs.update(resume_pkl=resume,
resume_kimg=resume_kimg,
resume_with_new_nets=True)
if ema_kimg is not None:
kwargs.update(G_ema_kimg=ema_kimg)
if d_aug:
kwargs.update(loss_args=loss_args)
else:
kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = train_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(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(data, train_dir, config, d_aug, diffaug_policy, cond, ops, mirror, mirror_v, \
kimg, batch_size, lrate, resume, resume_kimg, num_gpus, ema_kimg, gamma, freezeD):
# training functions
if d_aug: # https://github.com/mit-han-lab/data-efficient-gans
train = EasyDict(
run_func_name='training.training_loop_diffaug.training_loop'
) # Options for training loop (Diff Augment method)
loss_args = EasyDict(
func_name='training.loss_diffaug.ns_DiffAugment_r1',
policy=diffaug_policy) # Options for loss (Diff Augment method)
else: # original nvidia
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop (original from NVidia)
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
# network functions
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='1080p',
layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
G.impl = D.impl = ops
# dataset (tfrecords) - get or create
tfr_files = file_list(os.path.dirname(data), 'tfr')
tfr_files = [
f for f in tfr_files if basename(data) == basename(f).split('-')[0]
]
if len(tfr_files) == 0 or os.stat(tfr_files[0]).st_size == 0:
tfr_file, total_samples = create_from_image_folders(
data) if cond is True else create_from_images(data)
else:
tfr_file = tfr_files[0]
dataset_args = EasyDict(tfrecord=tfr_file)
# resolutions
with tf.Graph().as_default(), tflib.create_session().as_default(): # pylint: disable=not-context-manager
dataset_obj = dataset.load_dataset(
**dataset_args) # loading the data to see what comes out
resolution = dataset_obj.resolution
init_res = dataset_obj.init_res
res_log2 = dataset_obj.res_log2
dataset_obj.close()
dataset_obj = None
if list(init_res) == [4, 4]:
desc = '%s-%d' % (basename(data), resolution)
else:
print(' custom init resolution', init_res)
desc = basename(tfr_file)
G.init_res = D.init_res = list(init_res)
train.savenames = [desc.replace(basename(data), 'snapshot'), desc]
desc += '-%s' % config
# training schedule
train.total_kimg = kimg
train.image_snapshot_ticks = 1 * num_gpus if kimg <= 1000 else 4 * num_gpus
train.network_snapshot_ticks = 5
train.mirror_augment = mirror
train.mirror_augment_v = mirror_v
sched.tick_kimg_base = 2 if train.total_kimg < 2000 else 4
# learning rate
if config == 'e':
sched.G_lrate_base = 0.001
sched.G_lrate_dict = {0: 0.001, 1: 0.0007, 2: 0.0005, 3: 0.0003}
sched.lrate_step = 1500 # period for stepping to next lrate, in kimg
if config == 'f':
sched.G_lrate_base = lrate # 0.001 for big datasets, 0.0003 for few-shot
sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway
# batch size (for 16gb memory GPU)
sched.minibatch_gpu_base = 4096 // resolution if batch_size is None else batch_size
print(' Batch size', sched.minibatch_gpu_base)
sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base
sc.num_gpus = num_gpus
if config == 'e':
G.fmap_base = D.fmap_base = 8 << 10
if d_aug: loss_args.gamma = 100 if gamma is None else gamma
else: D_loss.gamma = 100 if gamma is None else gamma
elif config == 'f':
G.fmap_base = D.fmap_base = 16 << 10
else:
print(' Only configs E and F are implemented')
exit()
if cond:
desc += '-cond'
dataset_args.max_label_size = 'full' # conditioned on full label
if freezeD:
D.freezeD = True
train.resume_with_new_nets = True
if d_aug:
desc += '-daug'
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
tf_config=tf_config)
kwargs.update(resume_pkl=resume,
resume_kimg=resume_kimg,
resume_with_new_nets=True)
if ema_kimg is not None:
kwargs.update(G_ema_kimg=ema_kimg)
if d_aug:
kwargs.update(loss_args=loss_args)
else:
kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = train_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma,
mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='8k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 10
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
desc += '-' + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id not in ['config-f', 'config-l']:
G.fmap_base = D.fmap_base = 8 << 10
# Config L: Generator training only
if config_id == 'config-l':
# Use labels as latent vector input
dataset_args.max_label_size = "full"
# Deactivate methods specific for GAN training
G.truncation_psi = None
G.randomize_noise = False
G.style_mixing_prob = None
G.dlatent_avg_beta = None
G.conditional_labels = False
# Refinement training
G_loss.func_name = 'training.loss.G_reconstruction'
train.run_func_name = 'training.training_loop.training_loop_refinement'
# G.freeze_layers = ["mapping", "noise"]#, "4x4", "8x8", "16x16", "32x32"]
# Network for refinement
train.resume_pkl = "nets/stylegan2-ffhq-config-f.pkl" # TODO init net
train.resume_with_new_nets = True
# Maintenance tasks
sched.tick_kimg_base = 1 # 1 tick = 5000 images (metric update)
sched.tick_kimg_dict = {}
train.image_snapshot_ticks = 5 # Save every 5000 images
train.network_snapshot_ticks = 10 # Save every 10000 images
# Training parameters
sched.G_lrate_base = 1e-4
train.G_smoothing_kimg = 0.0
sched.minibatch_size_base = sched.minibatch_gpu_base * num_gpus # 4 per GPU
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {
128: 0.0015,
256: 0.002,
512: 0.003,
1024: 0.003
}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict(
tfrecord_dir='streetviewtf') # 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.
metrics = [] # 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.resume_snapshot = '/root/repos/idinvert/results/00016-sgan-1gpu/network-snapshot-015040.pkl' # Network pickle to resume training from, None = train from scratch.
train.resume_kimg = 15000 # Assumed training progress at the beginning. Affects reporting and training schedule.
train.resume_time = 0.0 # Assumed wallclock time at the beginning. Affects reporting.
train.mirror_augment = True
train.image_snapshot_ticks = 1
train.network_snapshot_ticks = 2
# Dataset.
#desc += '-ffhq'; dataset = EasyDict(tfrecord_dir='ffhq'); train.mirror_augment = True
##desc += '-ffhq512'; dataset = EasyDict(tfrecord_dir='ffhq', resolution=512); train.mirror_augment = True
#desc += '-ffhq256'; dataset = EasyDict(tfrecord_dir='ffhq', resolution=256); train.mirror_augment = True
#desc += '-celebahq'; dataset = EasyDict(tfrecord_dir='celebahq'); train.mirror_augment = True
#desc += '-bedroom'; dataset = EasyDict(tfrecord_dir='lsun-bedroom-full'); train.mirror_augment = False
#desc += '-car'; dataset = EasyDict(tfrecord_dir='lsun-car-512x384'); train.mirror_augment = False
#desc += '-cat'; dataset = EasyDict(tfrecord_dir='lsun-cat-full'); train.mirror_augment = False
# Number of GPUs.
desc += '-1gpu'
submit_config.num_gpus = 1
sched.minibatch_base = 16 #; sched.minibatch_dict = {4: 128, 8: 128, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8, 512: 4}
32: 64,
64: 32,
128: 16,
256: 8,
512: 4
}
#desc += '-2gpu'; submit_config.num_gpus = 2; sched.minibatch_base = 8; sched.minibatch_dict = {4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8}
#desc += '-4gpu'; submit_config.num_gpus = 4; sched.minibatch_base = 16; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16}
#desc += '-8gpu'; submit_config.num_gpus = 8; sched.minibatch_base = 32; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32}
# Default options.
train.total_kimg = args.total_kimg
sched.lod_initial_resolution = args.init_res
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)
train.image_snapshot_ticks = args.nsteps_image_snapshot
train.network_snapshot_ticks = args.nsteps_network_snapshot
train.resume_run_id = args.resume_run_id
train.resume_snapshot = args.resume_snapshot
# WGAN-GP loss for CelebA-HQ.
#desc += '-wgangp'; G_loss = EasyDict(func_name='training.loss.G_wgan'); D_loss = EasyDict(func_name='training.loss.D_wgan_gp'); sched.G_lrate_dict = {k: min(v, 0.002) for k, v in sched.G_lrate_dict.items()}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)
# Table 1.
#desc += '-tuned-baseline'; G.use_styles = False; G.use_pixel_norm = True; G.use_instance_norm = False; G.mapping_layers = 0; G.truncation_psi = None; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False
#desc += '-add-mapping-and-styles'; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False
#desc += '-remove-traditional-input'; G.style_mixing_prob = 0.0; G.use_noise = False
#desc += '-add-noise-inputs'; G.style_mixing_prob = 0.0
#desc += '-mixing-regularization' # default
# Table 2.
