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)
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
}
# Default options.
train.total_kimg = 20000
sched.lod_initial_resolution = 4
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)
desc += '-cond'
dataset.max_label_size = 'full'
# ----------------------------------------------------------------------------
# Main entry point for training.
# Calls the function indicated by 'train' using the selected options.
def main():
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
# Dataset.
desc += '-ffhq'; dataset = EasyDict(tfrecord_dir='stylegan'); 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
# Config presets from Progressive GAN (choose one).
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}; sched.G_lrate_dict = {1024: 0.0015}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000
#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}; sched.G_lrate_dict = {512: 0.0015, 1024: 0.002}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000
#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}; sched.G_lrate_dict = {256: 0.0015, 512: 0.002, 1024: 0.003}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000
#desc += '-8gpu'; submit_config.num_gpus = 8; sched.minibatch_base = 32; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32}; 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.total_kimg = 12000
# Tuned config for StyleGAN.
train.total_kimg = 25000; sched.lod_initial_resolution = 8
# 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.
#desc += '-mix0'; G.style_mixing_prob = 0.0
#desc += '-mix50'; G.style_mixing_prob = 0.5
#desc += '-mix90'; G.style_mixing_prob = 0.9 # default
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 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 run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.training_loop_mirror_v6_remove_half_fl_fr.training_loop')
G = EasyDict(func_name='training.networks.networks_stylegan2.G_main')
D = EasyDict(func_name='training.networks.networks_stylegan2_discriminator_new_rotation.D_stylegan2_new_rotaion') # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
G_loss = EasyDict(func_name='training.loss.loss_G_new_rotation_squared_euclidean_10_interpolate_50_percent_uniform_dist_int_penalty.G_logistic_ns_pathreg')
D_loss = EasyDict(func_name='training.loss.loss_D_logistic_r1_new_rotation_euclidean_square.D_logistic_r1_new_rotation')
sched = EasyDict()
grid = EasyDict(size='1080p', layout='random')
sc = dnnlib.SubmitConfig()
tf_config = {'rnd.np_random_seed': 1000}
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 10
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
# train.resume_pkl = './results/00200-stylegan2-car_labels_v7_oversample_filter-2gpu-config-f-squared_euclidean_10_interpolate_50_percent_int_reg-256/network-snapshot-000887.pkl'
# train.resume_kimg = 887.2
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
G.style_mixing_prob = None
desc += '-' + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
desc += '-squared_euclidean_10_interpolate_50_percent_int_reg_remove_half_fl_fr_no_noise_square'
desc += '-256'
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id != 'config-f':
G.fmap_base = D.fmap_base = 8 << 10
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(dataset,
data_dir,
result_dir,
config_id,
num_gpus,
total_kimg,
gamma,
mirror_augment,
metrics,
resume_run_id=None):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='8k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
if resume_run_id is not None:
# Resume from the ID of the results directory given
ids = sorted(get_valid_runids(result_dir))
if resume_run_id == 'recent':
resume_run_id = ids[-1][0]
else:
try:
resume_run_id = int(resume_run_id)
except ValueError:
raise RuntimeError(
'--resume argument is invalid (must be number, or "recent"): {}'
.format(resume_run_id))
try:
rundir_name = next(x[1] for x in ids if x[0] == resume_run_id)
except StopIteration:
raise RuntimeError(
'Could not find results directory with run ID {} (options: {})'
.format(resume_run_id, [x[0] for x in ids]))
# Find kimg & pkl file
rundir = os.path.join(result_dir, rundir_name)
pkls = [
name for name in os.listdir(rundir)
if name.startswith('network-snapshot-') and name.endswith('.pkl')
]
kimgs = sorted([(int(
pkl.replace('network-snapshot-', '').replace('.pkl', '')), pkl)
for pkl in pkls],
key=lambda x: x[0])
if len(kimgs) == 0:
raise RuntimeError(
'No network-snapshot-[0-9].pkl files found in {}'.format(
rundir))
max_kimg = kimgs[-1][0]
pkl_name = kimgs[-1][1]
# Get wall clock time
logfilepath = os.path.join(rundir, 'log.txt')
with open(logfilepath, 'r') as f:
logfile = f.read()
for line in logfile.splitlines():
if 'kimg {}'.format(max_kimg) in line:
if 'time ' not in line:
raise RuntimeError(
'Invalid log file: {}'.format(logfilepath))
line = line.split('time ')[1]
if 'sec/tick' not in line:
raise RuntimeError(
'Invalid log file: {}'.format(logfilepath))
line = line.split('sec/tick')[0].strip()
# Parse d h m s, etc.
total_seconds_formatted = line
total_seconds = 0
if 'd' in line:
arr = line.split('d')
days = int(arr[0].strip())
total_seconds += days * 24 * 60 * 60
line = arr[1]
if 'h' in line:
arr = line.split('h')
hours = int(arr[0].strip())
total_seconds += hours * 60 * 60
line = arr[1]
if 'm' in line:
arr = line.split('m')
mins = int(arr[0].strip())
total_seconds += mins * 60
line = arr[1]
if 's' in line:
arr = line.split('s')
secs = int(arr[0].strip())
total_seconds += secs
line = arr[1]
break
# Set args for training
train.resume_pkl = os.path.join(rundir, pkl_name)
train.resume_kimg = max_kimg
train.resume_time = total_seconds
print('Resuming from run {}: kimg {}, time {}'.format(
rundir_name, max_kimg, total_seconds_formatted))
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 1
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
desc += '-' + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id != 'config-f':
G.fmap_base = D.fmap_base = 8 << 10
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {
128: 0.0015,
256: 0.002,
512: 0.003,
1024: 0.003
}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
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_kimg = 120 # Should be enough in most cases.
train.total_kimg = train_kimg
train.resume_run_id = "models/2019-02-26-stylegan-faces-network-02048-016041.pkl"
train.network_snapshot_ticks = 1 # Save more frequently.
sched.lod_initial_resolution = 512 # Train from the highest resolution.
sched.lod_training_kimg = train_kimg
sched.lod_transition_kimg = train_kimg
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)
sched.tick_kimg_dict = {
4: 160,
8: 140,
16: 120,
32: 100,
64: 80,
128: 60,
256: 40,
#512:30,
512: 2, # 2k img / tick
1024: 20
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma,
mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='8k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 10
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
desc += '-' + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id not in ['config-f', 'config-l']:
G.fmap_base = D.fmap_base = 8 << 10
# Config L: Generator training only
if config_id == 'config-l':
# Use labels as latent vector input
dataset_args.max_label_size = "full"
# Deactivate methods specific for GAN training
G.truncation_psi = None
G.randomize_noise = False
G.style_mixing_prob = None
G.dlatent_avg_beta = None
G.conditional_labels = False
# Refinement training
G_loss.func_name = 'training.loss.G_reconstruction'
train.run_func_name = 'training.training_loop.training_loop_refinement'
# G.freeze_layers = ["mapping", "noise"]#, "4x4", "8x8", "16x16", "32x32"]
# Network for refinement
train.resume_pkl = "nets/stylegan2-ffhq-config-f.pkl" # TODO init net
train.resume_with_new_nets = True
# Maintenance tasks
sched.tick_kimg_base = 1 # 1 tick = 5000 images (metric update)
sched.tick_kimg_dict = {}
train.image_snapshot_ticks = 5 # Save every 5000 images
train.network_snapshot_ticks = 10 # Save every 10000 images
# Training parameters
sched.G_lrate_base = 1e-4
train.G_smoothing_kimg = 0.0
sched.minibatch_size_base = sched.minibatch_gpu_base * num_gpus # 4 per GPU
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {
128: 0.0015,
256: 0.002,
512: 0.003,
1024: 0.003
}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def main(args):
desc = 'sgan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict(tfrecord_dir=args.tfrecord_dirname,
resolution=args.resolution)
desc += '-ffhq' + str(args.resolution)
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='4k', layout='random') # Options for setup_snapshot_image_grid().
metrics = [metric_base.fid50k] # Options for MetricGroup.
submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.mirror_augment = True
# Number of GPUs.
submit_config.num_gpus = args.n_gpus
desc += '-{}gpu'.format(args.n_gpus)
sched.minibatch_base = 4 * args.n_gpus
sched.minibatch_dict = MINI_BATCH_DICT[args.n_gpus]
# Default options.
train.total_kimg = args.total_kimg
sched.lod_initial_resolution = 8
sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)
# kwargs
kwargs = EasyDict(train)
kwargs.update(data_root_dir=args.datadir)
kwargs.update(resume_run_id=args.resume_run_id,
resume_snapshot=args.resume_snapshot,
resume_kimg=args.resume_kimg,
resume_time=args.resume_time)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.run_dir_root = args.outdir
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def main():
parser = ArgumentParser(description='Style GAN')
parser.add_argument('--algorithm',
metavar='algorithm',
type=str,
help='algorithm',
default='sgan')
parser.add_argument('--resume_run_id',
metavar='resume_run_id',
type=str,
help='resume run id',
default=None)
parser.add_argument('--resume_kimg',
metavar='resume_kimg',
type=float,
help='resume kimg',
default=0.0)
args = parser.parse_args()
# ----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
if args.algorithm == 'sgan':
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().
resume_run_id = args.resume_run_id
resume_kimg = args.resume_kimg
tf_config = {
'rnd.np_random_seed': 1000
} # Options for tflib.init_tf().
# 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
desc += '-anime-faces-64'
dataset = EasyDict(tfrecord_dir='anime-faces-64', resolution=64)
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}
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 = 25000
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)
# 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.
# desc += '-mix0'; G.style_mixing_prob = 0.0
# desc += '-mix50'; G.style_mixing_prob = 0.5
# desc += '-mix90'; G.style_mixing_prob = 0.9 # default
# desc += '-mix100'; G.style_mixing_prob = 1.0
# Table 4.
# desc += '-traditional-0'; 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 += '-traditional-8'; G.use_styles = False; G.use_pixel_norm = True; G.use_instance_norm = False; G.mapping_layers = 8; G.truncation_psi = None; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False
# desc += '-stylebased-0'; G.mapping_layers = 0
# desc += '-stylebased-1'; G.mapping_layers = 1
# desc += '-stylebased-2'; G.mapping_layers = 2
# desc += '-stylebased-8'; G.mapping_layers = 8 # default
# ----------------------------------------------------------------------------
# Official training configs for Progressive GAN, targeted mainly for CelebA-HQ.
if args.algorithm == 'pgan':
desc = 'pgan' # 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.
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().
resume_run_id = args.resume_run_id
resume_kimg = args.resume_kimg
tf_config = {
'rnd.np_random_seed': 1000
} # Options for tflib.init_tf().
# Dataset (choose one).
desc += '-celebahq'
dataset = EasyDict(tfrecord_dir='celebahq')
train.mirror_augment = True
# desc += '-celeba'; dataset = EasyDict(tfrecord_dir='celeba'); train.mirror_augment = True
# desc += '-cifar10'; dataset = EasyDict(tfrecord_dir='cifar10')
# desc += '-cifar100'; dataset = EasyDict(tfrecord_dir='cifar100')
# desc += '-svhn'; dataset = EasyDict(tfrecord_dir='svhn')
# desc += '-mnist'; dataset = EasyDict(tfrecord_dir='mnist')
# desc += '-mnistrgb'; dataset = EasyDict(tfrecord_dir='mnistrgb')
# desc += '-syn1024rgb'; dataset = EasyDict(class_name='training.dataset.SyntheticDataset', resolution=1024, num_channels=3)
# desc += '-lsun-airplane'; dataset = EasyDict(tfrecord_dir='lsun-airplane-100k'); train.mirror_augment = True
# desc += '-lsun-bedroom'; dataset = EasyDict(tfrecord_dir='lsun-bedroom-100k'); train.mirror_augment = True
# desc += '-lsun-bicycle'; dataset = EasyDict(tfrecord_dir='lsun-bicycle-100k'); train.mirror_augment = True
# desc += '-lsun-bird'; dataset = EasyDict(tfrecord_dir='lsun-bird-100k'); train.mirror_augment = True
# desc += '-lsun-boat'; dataset = EasyDict(tfrecord_dir='lsun-boat-100k'); train.mirror_augment = True
# desc += '-lsun-bottle'; dataset = EasyDict(tfrecord_dir='lsun-bottle-100k'); train.mirror_augment = True
# desc += '-lsun-bridge'; dataset = EasyDict(tfrecord_dir='lsun-bridge-100k'); train.mirror_augment = True
# desc += '-lsun-bus'; dataset = EasyDict(tfrecord_dir='lsun-bus-100k'); train.mirror_augment = True
# desc += '-lsun-car'; dataset = EasyDict(tfrecord_dir='lsun-car-100k'); train.mirror_augment = True
# desc += '-lsun-cat'; dataset = EasyDict(tfrecord_dir='lsun-cat-100k'); train.mirror_augment = True
# desc += '-lsun-chair'; dataset = EasyDict(tfrecord_dir='lsun-chair-100k'); train.mirror_augment = True
# desc += '-lsun-churchoutdoor'; dataset = EasyDict(tfrecord_dir='lsun-churchoutdoor-100k'); train.mirror_augment = True
# desc += '-lsun-classroom'; dataset = EasyDict(tfrecord_dir='lsun-classroom-100k'); train.mirror_augment = True
# desc += '-lsun-conferenceroom'; dataset = EasyDict(tfrecord_dir='lsun-conferenceroom-100k'); train.mirror_augment = True
# desc += '-lsun-cow'; dataset = EasyDict(tfrecord_dir='lsun-cow-100k'); train.mirror_augment = True
# desc += '-lsun-diningroom'; dataset = EasyDict(tfrecord_dir='lsun-diningroom-100k'); train.mirror_augment = True
# desc += '-lsun-diningtable'; dataset = EasyDict(tfrecord_dir='lsun-diningtable-100k'); train.mirror_augment = True
# desc += '-lsun-dog'; dataset = EasyDict(tfrecord_dir='lsun-dog-100k'); train.mirror_augment = True
# desc += '-lsun-horse'; dataset = EasyDict(tfrecord_dir='lsun-horse-100k'); train.mirror_augment = True
# desc += '-lsun-kitchen'; dataset = EasyDict(tfrecord_dir='lsun-kitchen-100k'); train.mirror_augment = True
# desc += '-lsun-livingroom'; dataset = EasyDict(tfrecord_dir='lsun-livingroom-100k'); train.mirror_augment = True
# desc += '-lsun-motorbike'; dataset = EasyDict(tfrecord_dir='lsun-motorbike-100k'); train.mirror_augment = True
# desc += '-lsun-person'; dataset = EasyDict(tfrecord_dir='lsun-person-100k'); train.mirror_augment = True
# desc += '-lsun-pottedplant'; dataset = EasyDict(tfrecord_dir='lsun-pottedplant-100k'); train.mirror_augment = True
# desc += '-lsun-restaurant'; dataset = EasyDict(tfrecord_dir='lsun-restaurant-100k'); train.mirror_augment = True
# desc += '-lsun-sheep'; dataset = EasyDict(tfrecord_dir='lsun-sheep-100k'); train.mirror_augment = True
# desc += '-lsun-sofa'; dataset = EasyDict(tfrecord_dir='lsun-sofa-100k'); train.mirror_augment = True
# desc += '-lsun-tower'; dataset = EasyDict(tfrecord_dir='lsun-tower-100k'); train.mirror_augment = True
# desc += '-lsun-train'; dataset = EasyDict(tfrecord_dir='lsun-train-100k'); train.mirror_augment = True
# desc += '-lsun-tvmonitor'; dataset = EasyDict(tfrecord_dir='lsun-tvmonitor-100k'); train.mirror_augment = True
# Conditioning & snapshot options.
# desc += '-cond'; dataset.max_label_size = 'full' # conditioned on full label
# desc += '-cond1'; dataset.max_label_size = 1 # conditioned on first component of the label
# desc += '-g4k'; grid.size = '4k'
# desc += '-grpc'; grid.layout = 'row_per_class'
# Config presets (choose one).
# desc += '-preset-v1-1gpu'; submit_config.num_gpus = 1; D.mbstd_group_size = 16; sched.minibatch_base = 16; sched.minibatch_dict = {256: 14, 512: 6, 1024: 3}; sched.lod_training_kimg = 800; sched.lod_transition_kimg = 800; train.total_kimg = 19000
desc += '-preset-v2-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
}
sched.G_lrate_dict = {
1024: 0.0015
}
sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)
train.total_kimg = 12000
# desc += '-preset-v2-2gpus'; 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}; sched.G_lrate_dict = {512: 0.0015, 1024: 0.002}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000
# desc += '-preset-v2-4gpus'; submit_config.num_gpus = 4; sched.minibatch_base = 16; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16}; sched.G_lrate_dict = {256: 0.0015, 512: 0.002, 1024: 0.003}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000
# desc += '-preset-v2-8gpus'; submit_config.num_gpus = 8; sched.minibatch_base = 32; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32}; 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.total_kimg = 12000
# Numerical precision (choose one).
desc += '-fp32'
sched.max_minibatch_per_gpu = {
256: 16,
512: 8,
1024: 4
}
# desc += '-fp16'; G.dtype = 'float16'; D.dtype = 'float16'; G.pixelnorm_epsilon=1e-4; G_opt.use_loss_scaling = True; D_opt.use_loss_scaling = True; sched.max_minibatch_per_gpu = {512: 16, 1024: 8}
# Disable individual features.
# desc += '-nogrowing'; sched.lod_initial_resolution = 1024; sched.lod_training_kimg = 0; sched.lod_transition_kimg = 0; train.total_kimg = 10000
# desc += '-nopixelnorm'; G.use_pixelnorm = False
# desc += '-nowscale'; G.use_wscale = False; D.use_wscale = False
# desc += '-noleakyrelu'; G.use_leakyrelu = False
# desc += '-nosmoothing'; train.G_smoothing_kimg = 0.0
# desc += '-norepeat'; train.minibatch_repeats = 1
# desc += '-noreset'; train.reset_opt_for_new_lod = False
# Special modes.
# desc += '-BENCHMARK'; sched.lod_initial_resolution = 4; sched.lod_training_kimg = 3; sched.lod_transition_kimg = 3; train.total_kimg = (8*2+1)*3; sched.tick_kimg_base = 1; sched.tick_kimg_dict = {}; train.image_snapshot_ticks = 1000; train.network_snapshot_ticks = 1000
# desc += '-BENCHMARK0'; sched.lod_initial_resolution = 1024; train.total_kimg = 10; sched.tick_kimg_base = 1; sched.tick_kimg_dict = {}; train.image_snapshot_ticks = 1000; train.network_snapshot_ticks = 1000
# desc += '-VERBOSE'; sched.tick_kimg_base = 1; sched.tick_kimg_dict = {}; train.image_snapshot_ticks = 1; train.network_snapshot_ticks = 100
# desc += '-GRAPH'; train.save_tf_graph = True
# desc += '-HIST'; train.save_weight_histograms = True
# ----------------------------------------------------------------------------
# Main entry point for training.
# Calls the function indicated by 'train' using the selected options.
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.update(resume_run_id=resume_run_id, resume_kimg=resume_kimg)
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,
data_dir,
result_dir,
config_id,
num_gpus,
total_kimg,
gamma,
mirror_augment,
metrics,
resume_pkl=None,
resume_kimg=None,
):
train = EasyDict(
run_func_name="training.training_loop.training_loop",
# training resume options:
resume_pkl=
resume_pkl, # Network pickle to resume training from, None = train from scratch.
resume_kimg=
resume_kimg, # Assumed training progress at the beginning. Affects reporting and training schedule.
) # Options for training loop.
G = EasyDict(func_name="training.networks_stylegan2.G_main"
) # Options for generator network.
D = EasyDict(func_name="training.networks_stylegan2.D_stylegan2"
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name="training.loss.G_logistic_ns_pathreg"
) # Options for generator loss.
D_loss = EasyDict(func_name="training.loss.D_logistic_r1"
) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size="8k", layout="random") # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {"rnd.np_random_seed": 1000} # Options for tflib.init_tf().
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 10
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = "stylegan2"
desc += "-" + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += "-%dgpu" % num_gpus
assert config_id in _valid_configs
desc += "-" + config_id
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id != "config-f":
G.fmap_base = D.fmap_base = 8 << 10
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith("config-e"):
D_loss.gamma = 100
if "Gorig" in config_id:
G.architecture = "orig"
if "Gskip" in config_id:
G.architecture = "skip" # (default)
if "Gresnet" in config_id:
G.architecture = "resnet"
if "Dorig" in config_id:
D.architecture = "orig"
if "Dskip" in config_id:
D.architecture = "skip"
if "Dresnet" in config_id:
D.architecture = "resnet" # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ["config-a", "config-b", "config-c", "config-d"]:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {
128: 0.0015,
256: 0.002,
512: 0.003,
1024: 0.003,
}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = "G_synthesis_stylegan_revised"
D.func_name = "training.networks_stylegan2.D_stylegan"
# Configs A-C: Disable path length regularization.
if config_id in ["config-a", "config-b", "config-c"]:
G_loss = EasyDict(func_name="training.loss.G_logistic_ns")
# Configs A-B: Disable lazy regularization.
if config_id in ["config-a", "config-b"]:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == "config-a":
G = EasyDict(func_name="training.networks_stylegan.G_style")
D = EasyDict(func_name="training.networks_stylegan.D_basic")
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(
G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss,
)
kwargs.update(
dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config,
)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def initialize(self):
parser = self.parser
parser.add_argument('--config_file', type=argparse.FileType(mode='r'), help="configuration yml file")
self.parser.add_argument('--overwrite_config', action='store_true', help="overwrite config files if they exist")
self.parser.add_argument('--model', default='biggan', help="pretrained model to use, e.g. biggan, stylegan")
parser.add_argument('--transform', default="zoom", help="transform operation, e.g. zoom, shiftx, color, rotate2d"),
parser.add_argument('--num_samples', type=int, default=20000, help='number of latent z samples')
parser.add_argument('--loss', type=str, default='l2', help='loss to use for training', choices=['l2', 'lpips'])
parser.add_argument('--learning_rate', type=float, default=0.0001, help='learning rate for training')
parser.add_argument('--walk_type', type=str, default='NNz', choices=['NNz', 'linear'], help='type of latent z walk')
parser.add_argument('--models_dir', type=str, default="./models", help="output directory for saved checkpoints")
parser.add_argument('--model_save_freq', type=int, default=400, help="saves checkpoints after this many batches")
parser.add_argument('--name', type=str, help="experiment name, saved within models_dir")
parser.add_argument('--suffix', type=str, help="suffix for experiment name")
parser.add_argument('--prefix', type=str, help="prefix for experiment name")
parser.add_argument("--gpu", default="", type=str, help='GPUs to use (leave blank for CPU only)')
# NN walk parameters
group = parser.add_argument_group('nn', 'parameters used to specify NN walk')
group.add_argument('--eps', type=float, help="step size of each NN block")
group.add_argument('--num_steps', type=int, help="number of NN blocks")
# color transformation parameters
group = parser.add_argument_group('color', 'parameters used for color walk')
group.add_argument('--channel', type=int, help="which channel to modify; if unspecified, modifies all channels for linear walks, and luminance for NN walks")
# biggan walk parameters
group = parser.add_argument_group('biggan', 'parameters used for biggan walk')
group.add_argument('--category', type=int, help="which category to train on; if unspecified uses all categories")
# stylegan walk parameters
# Official training configs for StyleGAN, targeted mainly for anime.
if 1:
desc = 'sgan' # 包含在结果子目录名称中的描述字符串。
train = EasyDict(run_func_name='train.joint_train') # 训练过程设置。
G = EasyDict(func_name='training.networks_stylegan.G_style') # 生成网络架构设置。
D = EasyDict(func_name='training.networks_stylegan.D_basic') # 判别网络架构设置。
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.G_logistic_nonsaturating_steer') # 生成损失设置。
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp', r1_gamma=10.0) # 判别损失设置。
dataset = EasyDict() # 数据集设置，在后文确认。
sched = EasyDict() # 训练计划设置，在后文确认。
grid = EasyDict(size='4k', layout='random') # setup_snapshot_image_grid()相关设置。
metrics = [metric_base.fid50k] # 指标方法设置。
submit_config = dnnlib.SubmitConfig() # dnnlib.submit_run()相关设置。
tf_config = {'rnd.np_random_seed': 1000} # tflib.init_tf()相关设置。
# 数据集。
desc += '-character'; dataset = EasyDict(tfrecord_dir='character'); #train.mirror_augment = True
# GPU数量。
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}
#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}
# 默认设置。
train = EasyDict(total_kimg = 25000)
sched.lod_initial_resolution = 8
sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)
kwargs = EasyDict(is_train=True)
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
else:
kwargs = EasyDict(is_train=False)
group = parser.add_argument_group('stylegan', 'parameters used for stylegan walk')
group.add_argument('--dataset', default="anime", help="which dataset to use for pretrained stylegan, e.g. cars, cats, celebahq")
group.add_argument('--latent', default="w", help="which latent space to use; z or w")
group.add_argument('--truncation_psi', default=1.0, help="truncation for NN walk in w")
group.add_argument('--train_args', default=kwargs, help="kwargs for training stylegan")
# pgan walk parameters
group = parser.add_argument_group('pgan', 'parameters used for pgan walk')
group.add_argument('--dset', default="celebahq", help="which dataset to use for pretrained pgan")
self.initialized = True
return self.parser
def main():
#
# Usage example:
# python [this_file].py --kimg ###### --dataset [your data] --gpu_num 1
# --start_res 8 --img_size 512 --progressive True
#
#
# parse arguments
args = parse_args()
if args is None:
exit()
checkpoint_dir = args.checkpoint_dir
nvlabs_stylegan_pkl_kimg = args.kimg
nvlabs_stylegan_pkl_name = "network-snapshot-"+nvlabs_stylegan_pkl_kimg+".pkl"
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
# this is a hack since the taki0112 expects a dataset folder which may not exist
dataset = args.dataset
dataset_dir = "./dataset/" + dataset
temp_dataset_file = make_temp_dataset_file(dataset_dir)
# build the taki0112 StyleGAN architecture (vanilla Tensorflow)
gan = StyleGAN(sess, args)
# you have to go through this process to initialize everything needed to load the checkpoint...
original_start_res = args.start_res
args.start_res = args.img_size
gan.start_res = args.img_size
gan.build_model()
args.start_res = original_start_res
gan.start_res = original_start_res
# remove the temp file and the directory if it is empty
delete_temp_dataset_file(args, dataset_dir, temp_dataset_file)
# Initialize TensorFlow.
tflib.init_tf()
tf.global_variables_initializer().run()
vars = tf.trainable_variables("discriminator")
vars_vals = sess.run(vars)
for var, val in zip(vars, vars_vals):
print(var.name)
gan.saver = tf.train.Saver(max_to_keep=10)
gan.load(checkpoint_dir)
#
#
# Make an NVlabs StyleGAN network (default initialization)
#
#
# StyleGAN initialization parameters and options, if you care to change them, do so here
desc = "sgan"
train = EasyDict(run_func_name="training.training_loop.training_loop")
G = EasyDict(func_name="training.networks_stylegan.G_style")
D = EasyDict(func_name="training.networks_stylegan.D_basic")
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.G_logistic_nonsaturating")
D_loss = EasyDict(func_name="training.loss.D_logistic_simplegp", r1_gamma=10.0)
dataset = EasyDict()
sched = EasyDict()
grid = EasyDict(size="4k", layout="random")
metrics = [metric_base.fid50k]
submit_config = dnnlib.SubmitConfig()
tf_config = {"rnd.np_random_seed": 1000}
drange_net = [-1,1]
G_smoothing_kimg = 10.0
# Dataset.
desc += "-"+args.dataset
dataset = EasyDict(tfrecord_dir=args.dataset)
train.mirror_augment = True
# Number of GPUs.
gpu_num = args.gpu_num
if gpu_num == 1:
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}
elif gpu_num == 2:
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}
elif gpu_num == 4:
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}
elif gpu_num == 8:
desc += "-8gpu"; submit_config.num_gpus = 8
sched.minibatch_base = 32
sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32}
else:
print("ERROR: invalid number of gpus:",gpu_num)
sys.exit(-1)
# Default options.
train.total_kimg = 0
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)
# Initialize dnnlib and TensorFlow.
# ctx = dnnlib.RunContext(submit_config, train)
tflib.init_tf(tf_config)
# Construct networks.
with tf.device('/gpu:0'):
print('Constructing networks...')
dataset_resolution = args.img_size
dataset_channels = 3 # fairly sure everyone is using 3 channels ... # training_set.shape[0],
dataset_label_size = 0 # training_set.label_size,
G = tflib.Network('G',
num_channels=dataset_channels,
resolution=dataset_resolution,
label_size=dataset_label_size,
**G)
D = tflib.Network('D',
num_channels=dataset_channels,
resolution=dataset_resolution,
label_size=dataset_label_size,
**D)
Gs = G.clone('Gs')
G.print_layers(); D.print_layers()
print('Building TensorFlow graph...')
with tf.name_scope('Inputs'), tf.device('/cpu:0'):
lod_in = tf.placeholder(tf.float32, name='lod_in', shape=[])
lrate_in = tf.placeholder(tf.float32, name='lrate_in', shape=[])
minibatch_in = tf.placeholder(tf.int32, name='minibatch_in', shape=[])
minibatch_split = minibatch_in // submit_config.num_gpus
Gs_beta = 0.5 ** tf.div(tf.cast(minibatch_in, tf.float32),
G_smoothing_kimg * 1000.0) if G_smoothing_kimg > 0.0 else 0.0
src_d = "discriminator"
dst_d = "D"
src_gs = "generator/g_synthesis"
dst_gs = "G_synthesis" # "G_synthesis_1" <<<< this is handled later
src_gm = "generator/g_mapping"
dst_gm = "G_mapping" # "G_mapping_1" <<<< this is handled later
vars = tf.trainable_variables(src_gm)
vars_vals = sess.run(vars)
# Copy over the discriminator weights
for (new, old) in zip(tf.trainable_variables(dst_d), tf.trainable_variables(src_d)):
update_weight = [tf.assign(new, old)]
sess.run(update_weight)
temp_vals = sess.run([new, old])
# Copy over the Generator's mapping network weights
for (new, old) in zip(tf.trainable_variables(dst_gm), tf.trainable_variables(src_gm)):
update_weight = [tf.assign(new, old)]
sess.run(update_weight)
temp_vals = sess.run([new, old])
# Because the two network architectures use slightly different columns on one variable,
# you must set up code to handle the edge case transpose of the first case
first = True
for (new, old) in zip(tf.trainable_variables(dst_gs), tf.trainable_variables(src_gs)):
temp_vals = sess.run([new, old])
if new.shape != old.shape:
# you need a transpose with perm # old = tf.reshape(old, tf.shape(new))
# DO NOT USE RESHAPE
# (made this mistake here and the results work but are quite terrifying)
if (first):
first = False
old = tf.transpose(old, perm=[0, 3, 1, 2])
else:
old = tf.transpose(old, perm=[0, 1, 3, 2])
update_weight = [tf.assign(new, old)]
sess.run(update_weight)
# also update the running average network (not 100% sure this is necessary)
dst_gs = "G_synthesis_1"
dst_gm = "G_mapping_1"
for (new, old) in zip(tf.trainable_variables(dst_gm), tf.trainable_variables(src_gm)):
update_weight = [tf.assign(new, old)]
sess.run(update_weight)
temp_vals = sess.run([new, old])
first = True
for (new, old) in zip(tf.trainable_variables(dst_gs), tf.trainable_variables(src_gs)):
temp_vals = sess.run([new, old])
if new.shape != old.shape:
# you need a transpose with perm # old = tf.reshape(old, tf.shape(new))
# DO NOT USE RESHAPE
# (made this mistake here and the results work but are quite terrifying)
if (first):
first = False
old = tf.transpose(old, perm=[0, 3, 1, 2])
else:
old = tf.transpose(old, perm=[0, 1, 3, 2])
update_weight = [tf.assign(new, old)]
sess.run(update_weight)
# Also, assign the w_avg in the taki0112 network to the NVlabs Gs dlatent_avg
new = [x for x in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="G")
if "dlatent_avg" in str(x)][0] # G.get_var("dlatent_avg")
old = [x for x in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="generator")
if "avg" in str(x)][0]
update_weight = [tf.assign(new, old)]
sess.run(update_weight)
vars = [new]
vars_vals = gan.sess.run(vars)
vars_vals = sess.run(vars)
misc.save_pkl((G, D, Gs), "./"+nvlabs_stylegan_pkl_name)
4: 128,
8: 128,
16: 128,
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
