def prepare(tasks, data_dir, max_images = None,
ratio = 1.0, images_dir = None, format = None): # Options for custom dataset
os.makedirs(data_dir, exist_ok = True)
for task in tasks:
# If task not in catalog, create custom task configuration
c = catalog.get(task, EasyDict({
"local": True,
"name": task,
"dir": images_dir,
"ratio": ratio,
"process": formats_catalog.get(format)
}))
dirname = f"{data_dir}/{task}"
os.makedirs(dirname, exist_ok = True)
# try:
print(misc.bold(f"Preparing the {c.name} dataset..."))
if "local" not in c:
fname = f"{dirname}/{c.filename}"
download = not ((os.path.exists(fname) and verify_md5(fname, c.md5)))
path = get_path(c.url, dirname, path = c.filename)
if download:
print(misc.bold(f"Downloading the data ({c.size} GB)..."))
download_file(c.url, path)
if path.endswith(".zip"):
if not is_unzipped(path, dirname):
print(misc.bold(f"Unzipping {path}..."))
unzip(path, dirname)
if "process" in c:
imgdir = images_dir if "local" in c else (f"{dirname}/{c.dir}")
c.process(dirname, imgdir, ratio = c.ratio, max_imgs = max_images)
print(misc.bcolored(f"Completed preparations for {c.name}!", "blue"))
def main():
submit_config = dnnlib.SubmitConfig()
# Which metrics to evaluate?
metrics = []
metrics += [metric_base.fid50k]
#metrics += [metric_base.ppl_zfull]
#metrics += [metric_base.ppl_wfull]
#metrics += [metric_base.ppl_zend]
#metrics += [metric_base.ppl_wend]
#metrics += [metric_base.ls]
#metrics += [metric_base.dummy]
# Which networks to evaluate them on?
tasks = []
tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ', dataset_args=EasyDict(tfrecord_dir='ffhq', shuffle_mb=0), mirror_augment=True)] # karras2019stylegan-ffhq-1024x1024.pkl
#tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=100, snapshot=25000)]
#tasks += [EasyDict(run_func_name='run_metrics.run_all_snapshots', run_id=100)]
# How many GPUs to use?
submit_config.num_gpus = 1
#submit_config.num_gpus = 2
#submit_config.num_gpus = 4
#submit_config.num_gpus = 8
# Execute.
submit_config.run_dir_root = base_modules.ganspace.models.stylegan.stylegan_tf.dnnlib.submission.submit.get_template_from_path(
config.result_dir)
submit_config.run_dir_ignore += config.run_dir_ignore
for task in tasks:
for metric in metrics:
submit_config.run_desc = '%s-%s' % (task.run_func_name, metric.name)
if task.run_func_name.endswith('run_snapshot'):
submit_config.run_desc += '-%s-%s' % (task.run_id, task.snapshot)
if task.run_func_name.endswith('run_all_snapshots'):
submit_config.run_desc += '-%s' % task.run_id
submit_config.run_desc += '-%dgpu' % submit_config.num_gpus
dnnlib.submit_run(submit_config, metric_args=metric, **task)
def main():
kwargs = EasyDict(train)
kwargs.update(
G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss,
)
kwargs.update(
dataset_args=dataset,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config,
)
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(
config.result_dir)
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def main():
kwargs = EasyDict(train)
kwargs.update(Encoder_args=Encoder,
E_opt_args=E_opt,
D_opt_args=D_opt,
E_loss_args=E_loss,
D_loss_args=D_loss,
lr_args=lr)
kwargs.update(dataset_args=Data_dir,
decoder_pkl=Decoder_pkl,
tf_config=tf_config)
kwargs.lr_args.decay_step = train.max_iters // 4
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.num_gpus = num_gpus
kwargs.submit_config.image_size = image_size
kwargs.submit_config.batch_size = batch_size
kwargs.submit_config.batch_size_test = batch_size_test
kwargs.submit_config.z_dim = z_dim
kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(
config.result_dir)
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def easygen_train(model_path,
images_path,
dataset_path,
start_kimg=7000,
max_kimg=25000,
schedule='',
seed=1000):
#import stylegan
#from stylegan import config
##from stylegan import dnnlib
#from stylegan.dnnlib import EasyDict
#images_dir = '/content/raw'
#max_kimg = 25000
#start_kimg = 7000
#schedule = ''
#model_in = '/content/karras2019stylegan-cats-256x256.pkl'
#dataset_dir = '/content/stylegan_dataset' #os.path.join(cwd, 'cache', 'stylegan_dataset')
import config
config.data_dir = '/content/datasets'
config.results_dir = '/content/results'
config.cache_dir = '/contents/cache'
run_dir_ignore = [
'/contents/results', '/contents/datasets', 'contents/cache'
]
import copy
import dnnlib
from dnnlib import EasyDict
from metrics import metric_base
# Prep dataset
import dataset_tool
print("prepping dataset...")
dataset_tool.create_from_images(tfrecord_dir=dataset_path,
image_dir=images_path,
shuffle=False)
# Set up training parameters
desc = 'sgan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict() # Options for load_dataset().
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='1080p',
layout='random') # Options for setup_snapshot_image_grid().
#metrics = [metric_base.fid50k] # Options for MetricGroup.
submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': seed} # Options for tflib.init_tf().
# Dataset
desc += '-custom'
dataset = EasyDict(tfrecord_dir=dataset_path)
train.mirror_augment = True
# Number of GPUs.
desc += '-1gpu'
submit_config.num_gpus = 1
sched.minibatch_base = 4
sched.minibatch_dict = {
4: 128,
8: 128,
16: 128,
32: 64,
64: 32,
128: 16,
256: 8,
512: 4
} #{4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 16}
# Default options.
train.total_kimg = max_kimg
sched.lod_initial_resolution = 8
sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)
# schedule
schedule_dict = {
4: 160,
8: 140,
16: 120,
32: 100,
64: 80,
128: 60,
256: 40,
512: 30,
1024: 20
} #{4: 2, 8:2, 16:2, 32:2, 64:2, 128:2, 256:2, 512:2, 1024:2} # Runs faster for small datasets
if len(schedule) >= 5 and schedule[0] == '{' and schedule[
-1] == '}' and ':' in schedule:
# is schedule a string of a dict?
try:
temp = eval(schedule)
schedule_dict = dict(temp)
# assert: it is a dict
except:
pass
elif len(schedule) > 0:
# is schedule an int?
try:
schedule_int = int(schedule)
#assert: schedule is an int
schedule_dict = {}
for i in range(1, 10):
schedule_dict[int(math.pow(2, i + 1))] = schedule_int
except:
pass
print('schedule:', str(schedule_dict))
sched.tick_kimg_dict = schedule_dict
# resume kimg
resume_kimg = start_kimg
# path to model
resume_run_id = model_path
# tick snapshots
image_snapshot_ticks = 1
network_snapshot_ticks = 1
# Submit run
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset,
sched_args=sched,
grid_args=grid,
tf_config=tf_config)
kwargs.update(resume_kimg=resume_kimg, resume_run_id=resume_run_id)
kwargs.update(image_snapshot_ticks=image_snapshot_ticks,
network_snapshot_ticks=network_snapshot_ticks)
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(
config.result_dir)
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""Main entry point for training StyleGAN and ProGAN networks."""
import copy
import dnnlib
from dnnlib import EasyDict
import config
from metrics import metric_base
#----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
if 1:
desc = 'sgan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict() # Options for load_dataset().
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
import hashlib
import argparse
import numpy as np
from training import misc
from dnnlib import EasyDict
import dataset_tool
catalog = {
"ffhq":
EasyDict({
"name": "FFHQ", # Dataset name for logging
"filename": "ffhq-r08.tfrecords1of1", # Local file name
"url":
"http://downloads.cs.stanford.edu/nlp/data/dorarad/ffhq-r08.tfrecords1of1", # download URL
"md5":
"74de4f07dc7bfb07c0ad4471fdac5e67", # MD5 checksum to potentially skip download
"ratio": 1, # height/width ratio
"size": 13, # download size in GB
"shards": 1, # Number of tfrecord shards
"img_num": 70000 # Number of images
}),
"bedrooms":
EasyDict({
"name": "LSUN-Bedrooms", # Dataset name for logging
"filename": "bedroom_train_lmdb.zip",
"url": "http://dl.yf.io/lsun/scenes/bedroom_train_lmdb.zip",
"md5": "f2c5d904a82a6295dbdccb322b4b0a99",
"dir": "bedroom_train_lmdb",
"ratio": 188 / 256,
"size": 43,
"shards": 64,
def run(data_dir='./../datasets',
result_dir='./../results',
num_gpus=1,
total_kimg=1000,
mirror_augment=True):
train = EasyDict(
run_func_name='classifier_vgg.training_loop-vgg.training_loop')
classifier = EasyDict(
func_name='classifier_vgg.network_classifier-vgg.classifier_vgg')
classifier_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
classifier_loss = EasyDict(func_name='classifier_vgg.loss.cross_entropy')
sched = EasyDict()
sc = dnnlib.SubmitConfig()
tf_config = {'rnd.np_random_seed': 1000}
# train.resume_pkl = './results/00254-classifier-single_class_model/network-snapshot-001000.pkl'
# train.resume_kimg = 1000.0
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 8
current_label = 'model'
dataset = 'classifier_oversample_' + current_label
print(dataset)
train.data_dir = data_dir
if not os.path.exists(train.data_dir):
print('Error: dataset root directory does not exist.')
sys.exit(1)
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
metrics = [
EasyDict(name='acc_' + current_label,
func_name='metrics.accuracy_single_class.ACC',
test_dataset=dataset + '_test',
num_images=5000,
minibatch_per_gpu=8,
test_data_dir=data_dir)
]
desc = 'classifier-single_class_' + current_label
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(classifier_args=classifier,
classifier_opt_args=classifier_opt,
classifier_loss_args=classifier_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.rotation.v7_baseline.training_loop')
G = EasyDict(func_name='training.networks.rotation.v7_baseline.G_main')
D = EasyDict(func_name='training.networks.rotation.v7_baseline.D_stylegan2')
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
G_loss = EasyDict(func_name='training.loss.rotation.v7_int_reg_vgg.G_logistic_ns_pathreg')
D_loss = EasyDict(func_name='training.loss.rotation.v7_int_reg_vgg.D_logistic_r1')
sched = EasyDict()
grid = EasyDict(size='1080p', layout='random')
sc = dnnlib.SubmitConfig()
tf_config = {'rnd.np_random_seed': 1000}
# train.resume_pkl = './../results/00422-rotation_v7_int_reg_vgg/network-snapshot-003628.pkl'
# train.resume_kimg = 3628.0
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 10
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'rotation_v7_int_reg_vgg_increase'
G_loss.int_pl_weight = 20.0
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
assert config_id in _valid_configs
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""Main entry point for training StyleGAN networks."""
import copy
import dnnlib
from dnnlib import EasyDict
import argparse
import config
from metrics import metric_base
#----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
if 1:
desc = 'sgan'
train = EasyDict() # Description string included in result subdir name.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict() # Options for load_dataset().
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
import copy
import config
import dnnlib
from dnnlib import EasyDict
from metrics import metric_base
# ----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
if 1:
# Description string included in result subdir name.
desc = 'sgan'
# Options for training loop.
train = EasyDict(run_func_name='training.training_loop.training_loop')
# Options for generator network.
G = EasyDict(func_name='training.networks_stylegan.G_style')
# Options for discriminator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic')
# Options for generator optimizer.
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
# Options for discriminator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
# Options for generator loss.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating')
# Options for discriminator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0)
# Options for load_dataset().
dataset = EasyDict()
# Default metric definitions
from dnnlib import EasyDict
metric_defaults = EasyDict([(args.name, args) for args in [
EasyDict(name = "fid", func_name = "metrics.frechet_inception_distance.FID", minibatch_per_gpu = 8),
EasyDict(name = "is", func_name = "metrics.inception_score.IS", num_splits = 10, minibatch_per_gpu = 8),
EasyDict(name = "pr", func_name = "metrics.precision_recall.PR", nhood_size = 3, minibatch_per_gpu = 8, row_batch_size = 5000, col_batch_size = 5000),
# EasyDict(name = "ppl_zfull", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "z", sampling = "full", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")),
# EasyDict(name = "ppl_wfull", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "w", sampling = "full", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")),
# EasyDict(name = "ppl_zend", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "z", sampling = "end", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")),
# EasyDict(name = "ppl_wend", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "w", sampling = "end", crop = True, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")),
# EasyDict(name = "ppl2_wend", func_name = "metrics.perceptual_path_length.PPL", num_samples = 50000, epsilon = 1e-4, space = "w", sampling = "end", crop = False, minibatch_per_gpu = 4, Gs_overrides = dict(dtype = "float32", mapping_dtype = "float32")),
# EasyDict(name = "ls", func_name = "metrics.linear_separability.LS", num_samples = 200000, num_keep = 100000, attrib_indices = range(40), minibatch_per_gpu = 4),
]])
import copy
import os
import dnnlib
from dnnlib import EasyDict
import config
from metrics import metric_base
# ----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
# turn off black formatting for this file:
# fmt: off
desc = 'msg-stylegan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict() # Options for load_dataset().
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(size='4k',
import os
import dnnlib
from dnnlib import EasyDict
import config
import copy
desc = 'stylegan-encoder' + config.DESC
train = EasyDict(run_func_name='training.training_loop_encoder.training_loop')
Encoder = EasyDict(func_name='training.networks_encoder.Encoder')
E_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8)
D_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8)
E_loss = EasyDict(func_name='training.loss_encoder.E_loss',
feature_scale=0.00005,
D_scale=0.05,
perceptual_img_size=256)
D_loss = EasyDict(func_name='training.loss_encoder.D_logistic_simplegp',
r1_gamma=10.0)
lr = EasyDict(learning_rate=0.0001,
decay_step=30000,
decay_rate=0.8,
stair=False)
Data_dir = EasyDict(data_train=config.DATA_DIR, data_test=config.TEST_DATA_DIR)
Decoder_pkl = EasyDict(decoder_pkl=config.PICKLE_DIR)
Inversion_pkl = EasyDict(inversion_pkl=config.INVERSION_PICKLE_DIR)
tf_config = {'rnd.np_random_seed': 1000}
submit_config = dnnlib.SubmitConfig()
num_gpus = 1
desc += '-1gpu'
# num_gpus = 2; desc += '-2gpu'
# num_gpus = 4; desc += '-4gpu'
import os
import dnnlib
from dnnlib import EasyDict
import config
import copy
#os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
desc = 'LIA'
train = EasyDict(run_func_name='training.training_loop_encoder.training_loop')
Encoder = EasyDict(func_name='training.networks_encoder.Encoder')
E_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8)
D_opt = EasyDict(beta1=0.9, beta2=0.99, epsilon=1e-8)
E_loss = EasyDict(func_name='training.loss_encoder.E_perceptual_loss',
feature_scale=0.00005,
D_scale=0.1)
D_loss = EasyDict(func_name='training.loss_encoder.D_logistic_simplegp',
r1_gamma=10.0)
lr = EasyDict(learning_rate=0.0001,
decay_step=50000,
decay_rate=0.8,
stair=False)
Data_dir = EasyDict(data_train='../data/ffhq_128_train_6w5.tfrecords',
data_test='../data/ffhq_128_test_5k.tfrecords')
Decoder_pkl = EasyDict(decoder_pkl='./network-snapshot-023315.pkl'
) #the first stage training results
tf_config = {'rnd.np_random_seed': 1000}
submit_config = dnnlib.SubmitConfig()
# num_gpus = 1; desc += '-1gpu'
# num_gpus = 2; desc += '-2gpu'
"""Main entry point for training StyleGAN and ProGAN networks."""
import copy
import dnnlib
from dnnlib import EasyDict
import config
from metrics import metric_base
#----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
if 1:
desc = 'sgan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop') # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style') # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic') # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating') # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp', r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict() # Options for load_dataset().
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(size='4k', layout='random') # Options for setup_snapshot_image_grid().
metrics = [metric_base.fid50k] # Options for MetricGroup.
submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
# Dataset.
desc += '-images_512x512'; dataset = EasyDict(tfrecord_dir='images_512x512'); train.mirror_augment = True
def main():
kwargs = EasyDict(train)
kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss)
kwargs.update(dataset_args=dataset, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(config.result_dir)
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
kwargs.save_weight_histograms = True
kwargs.network_snapshot_ticks = 5
kwargs.update(resume_run_id=-1)
dnnlib.submit_run(**kwargs)
def create_model(config_id='config-f',
gamma=None,
height=512,
width=512,
cond=None,
label_size=0):
train = EasyDict(run_func_name='training.diagnostic.create_initial_pkl'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
sched.minibatch_size_base = 192
sched.minibatch_gpu_base = 3
D_loss.gamma = 10
desc = 'stylegan2'
dataset_args = EasyDict() # (tfrecord_dir=dataset)
if cond:
desc += '-cond'
dataset_args.max_label_size = 'full' # conditioned on full label
desc += '-' + config_id
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id != 'config-f':
G.fmap_base = D.fmap_base = 8 << 10
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {
128: 0.0015,
256: 0.002,
512: 0.003,
1024: 0.003
}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
G.update(resolution_h=height)
G.update(resolution_w=width)
D.update(resolution_h=height)
D.update(resolution_w=width)
sc.submit_target = dnnlib.SubmitTarget.DIAGNOSTIC
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
# [EDITED]
kwargs.update(G_args=G,
D_args=D,
tf_config=tf_config,
config_id=config_id,
resolution_h=height,
resolution_w=width,
label_size=label_size)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_desc = desc
dnnlib.submit_diagnostic(**kwargs)
return f'network-initial-config-f-{height}x{width}-{label_size}.pkl'
def main():
#------------------------------------------------------------------------
# Modified by Deng et al.
args = parse_args()
if args is None:
exit()
weight_args = EasyDict()
weight_args.update(w_id=args.w_id,
w_lm=args.w_lm,
w_gamma=args.w_gamma,
w_skin=args.w_skin,
w_exp_warp=args.w_exp_warp,
w_gamma_change=args.w_gamma_change)
train.update(run_func_name='training.training_loop.training_loop')
kwargs = EasyDict(train)
# stage 1: training with only imitative losses with 15000k images.
if args.stage == 1:
train_stage = EasyDict(
func_name='training.training_utils.training_stage1')
kwargs.update(total_kimg=15000)
# stage 2: training with imitative losses and contrastive losses.
else:
train_stage = EasyDict(
func_name='training.training_utils.training_stage2')
kwargs.update(resume_run_id=args.run_id,
resume_snapshot=args.snapshot,
resume_kimg=args.kimg)
kwargs.update(total_kimg=25000)
weight_args.update(w_lm=100)
kwargs.update(train_stage_args=train_stage)
kwargs.update(weight_args=weight_args, noise_dim=args.noise_dim)
#------------------------------------------------------------------------
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(
config.result_dir)
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def main():
parser = argparse.ArgumentParser(
description='Train ProgressiveGAN w/ Hessian Penalty')
# Model/Dataset Parameters:
parser.add_argument('--num_gpus',
type=int,
required=True,
help='Number of GPUs to train with')
parser.add_argument(
'--dataset',
type=str,
default='edges_and_shoes',
help='Name of TFRecords directory in datasets/ folder to train on.')
parser.add_argument(
'--resolution',
type=int,
default=128,
help='Maximum resolution (after progressive growth) to train at.')
parser.add_argument('--nz',
type=int,
default=12,
help='Number of components in G\'s latent space.')
parser.add_argument('--total_kimg',
type=int,
default=30000,
help='How long to train for')
parser.add_argument(
'--resume_exp',
type=int,
default=None,
help=
'If specified, resumes training from a snapshot in the results directory with job number resume_exp'
)
parser.add_argument(
'--resume_snapshot',
default='latest',
help=
'If using resume_exp, you can override this default to resume training from a specific checkpoint'
)
parser.add_argument('--seed',
type=int,
default=1000,
help='NumPy Random Seed')
# Hessian Penalty Parameters:
parser.add_argument(
'--hp_lambda',
type=float,
default=0.1,
help='Loss weighting of the Hessian Penalty regularization term. '
'When fine-tuning with the Hessian Penalty, this value '
'is the maximum weighting once the loss term is fully ramped-up.'
'Set to 0.0 to disable the Hessian Penalty.')
parser.add_argument(
'--epsilon',
type=float,
default=0.1,
help='The granularity of the finite differences approximation. '
'When changing this value from 0.1, you will likely need to change '
'hp_lambda as well to get optimal results.')
parser.add_argument(
'--num_rademacher_samples',
type=int,
default=2,
help='The number of Rademacher vectors to be sampled per-batch element '
'when estimating the Hessian Penalty. Must be >=2. Setting this '
'parameter larger can result in GPU out-of-memory for smaller GPUs!')
parser.add_argument(
'--layers_to_reg',
nargs='*',
default=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
help=
'Which activations to directly apply the Hessian Penalty to? For 128x128 '
'resolution, you can choose any layers 1-13. '
'(1=first fully-connected layer, 2-12=convolutional layers, 13=RGB activations)'
)
parser.add_argument(
'--warmup_kimg',
type=int,
default=10000,
help='Over how many kimg (=thousands of real images fed to D) should '
'the Hessian Penalty weight be linearly ramped-up? Setting this '
'value too low can sometimes result in training instability when fine-tuning.'
'The warmup will start at hp_start_kimg.')
parser.add_argument(
'--hp_start_kimg',
type=int,
default=0,
help='To train from scratch with the Hessian Penalty, set this to 0. '
'If >0, the Hessian Penalty will be phased-in starting at the '
'specified iteration. Alternatively, you can load '
'a checkpoint and simply set this to 0 to fine-tune a pretrained model.'
'Note that we haven\'t tested setting this > 0.')
# InfoGAN Parameters:
parser.add_argument(
'--infogan_lambda',
type=float,
default=0.0,
help=
'Loss weighting for InfoGAN loss on latent vectors. Set to 0.0 to disable InfoGAN loss.'
)
parser.add_argument(
'--infogan_nz',
type=int,
default=0,
help=
'Number of Z components to regularize with InfoGAN\'s mutual information loss. '
'For a (batch_size, nz) Z vector, you can access the regularized components with '
'Z[:, :infogan_nz].')
# Visualization and Metric Parameters:
parser.add_argument(
'--compute_interp_ticks',
type=int,
default=5,
help=
'Training ticks between computing disentanglement visualizations (interpolations).'
)
parser.add_argument(
'--compute_metrics_ticks',
type=int,
default=5,
help=
'Training ticks between computing FID (and optionally PPL; see next parameter).'
)
parser.add_argument('--metrics',
nargs='*',
default=['FID', 'PPL'],
help='Which metrics to compute during training.')
parser.add_argument(
'--dashboard_api',
type=str,
choices=['tensorboard', 'wandb'],
default='tensorboard',
help=
'Which training dashboard software to use for logging visualizations/metrics/losses/etc; '
'either TensorBoard (default) or WandB. If you choose to use WandB, you will first need '
'an account on wandb.com.')
parser.add_argument('--wandb_entity',
type=str,
default=None,
help='If using WandB, the entity for logging'
'(e.g., your username)')
opt = parser.parse_args()
opt = EasyDict(vars(opt))
verify_opt(opt)
run(opt)
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""Main entry point for training StyleGAN and ProGAN networks."""
import copy
import dnnlib
from dnnlib import EasyDict
import config
from metrics import metric_base
#----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
if 1:
desc = 'sgan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict() # Options for load_dataset().
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
def run(opt):
"""Sets-up all of the parameters necessary to start a ProgressiveGAN training job."""
desc = build_job_name(
opt) # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_progan.G_paper'
) # Options for generator network.
D = EasyDict(func_name='training.networks_progan.D_paper'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(
func_name='training.loss.G_wgan') # Options for generator loss.
D_loss = EasyDict(
func_name='training.loss.D_wgan_gp') # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='1080p',
layout='random') # Options for setup_snapshot_image_grid().
submit_config = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {
'rnd.np_random_seed': opt.seed
} # Options for tflib.init_tf().
metrics = [] # Metrics to run during training.
if 'FID' in opt.metrics:
metrics.append(metric_base.fid50k)
if 'PPL' in opt.metrics:
metrics.append(metric_base.ppl_zend_v2)
train.network_metric_ticks = opt.compute_metrics_ticks
train.interp_snapshot_ticks = opt.compute_interp_ticks
find_dataset(opt.dataset)
# Optionally resume from checkpoint:
if opt.resume_exp is not None:
results_dir = os.path.join(os.getcwd(), config.result_dir)
_resume_pkl = find_pkl(results_dir, opt.resume_exp,
opt.resume_snapshot)
train.resume_run_id = opt.resume_exp
train.resume_snapshot = _resume_pkl
train.resume_kimg = int(_resume_pkl.split('.pkl')[0][-6:])
if f'hessian_penalty_{opt.dataset}' not in _resume_pkl and opt.hp_lambda > 0:
print(
'When fine-tuning a job that was originally trained without the Hessian Penalty, '
'hp_start_kimg is relative to the kimg of the checkpoint being resumed from. '
'Hessian Penalty will be phased-in starting at absolute '
f'kimg={opt.hp_start_kimg + train.resume_kimg}.')
opt.hp_start_kimg += train.resume_kimg
# Set up dataset hyper-parameters:
dataset = EasyDict(tfrecord_dir=os.path.join(os.getcwd(), config.data_dir,
opt.dataset),
resolution=opt.resolution)
train.mirror_augment = False
# Set up network hyper-parameters:
G.latent_size = opt.nz
D.infogan_nz = opt.infogan_nz
G.infogan_lambda = opt.infogan_lambda
D.infogan_lambda = opt.infogan_lambda
# When computing the multi-layer Hessian Penalty, we retrieve intermediate activations by accessing the
# corresponding tensor's name. Below are the names of various activations in G that we can retrieve:
activation_type = 'norm'
progan_generator_layer_index_to_name = {
1: f'4x4/Dense/Post_{activation_type}',
2: f'4x4/Conv/Post_{activation_type}',
3: f'8x8/Conv0_up/Post_{activation_type}',
4: f'8x8/Conv1/Post_{activation_type}',
5: f'16x16/Conv0_up/Post_{activation_type}',
6: f'16x16/Conv1/Post_{activation_type}',
7: f'32x32/Conv0_up/Post_{activation_type}',
8: f'32x32/Conv1/Post_{activation_type}',
9: f'64x64/Conv0_up/Post_{activation_type}',
10: f'64x64/Conv1/Post_{activation_type}',
11: f'128x128/Conv0_up/Post_{activation_type}',
12: f'128x128/Conv1/Post_{activation_type}',
13: 'images_out' # final full-resolution RGB activation
}
# Convert from layer indices to layer names (which we'll need to compute the Hessian Penalty):
layers_to_reg = [
progan_generator_layer_index_to_name[layer_ix]
for layer_ix in sorted(opt.layers_to_reg)
]
# Store the Hessian Penalty parameters in their own dictionary:
HP = EasyDict(hp_lambda=opt.hp_lambda,
epsilon=opt.epsilon,
num_rademacher_samples=opt.num_rademacher_samples,
layers_to_reg=layers_to_reg,
warmup_nimg=opt.warmup_kimg * 1000,
hp_start_nimg=opt.hp_start_kimg * 1000)
# How long to train for (as measured by thousands of real images processed, not gradient steps):
train.total_kimg = opt.total_kimg
# We ran the original experiments using 4 GPUs per job. If using a different number,
# we try to scale batch sizes up or down accordingly in the for-loop below. Note that
# using other batch sizes is somewhat untested, though!
submit_config.num_gpus = opt.num_gpus
sched.minibatch_base = 32
sched.minibatch_dict = {
4: 2048,
8: 1024,
16: 512,
32: 256,
64: 128,
128: 96,
256: 32,
512: 16
}
for res, batch_size in sched.minibatch_dict.items():
sched.minibatch_dict[res] = int(batch_size * opt.num_gpus / 4)
# Set-up WandB if optionally using it instead of TensorBoard:
if opt.dashboard_api == 'wandb':
init_wandb(opt=opt,
name=desc,
group=opt.dataset,
entity=opt.wandb_entity)
# Start the training job:
kwargs = EasyDict(train)
kwargs.update(HP_args=HP,
G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(submit_config)
kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(
config.result_dir)
kwargs.submit_config.run_dir_ignore += config.run_dir_ignore
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(dataset, resolution, result_dir, DiffAugment, num_gpus, batch_size,
total_kimg, ema_kimg, num_samples, gamma, fmap_base, fmap_max,
latent_size, mirror_augment, impl, metrics, resume, resume_kimg,
num_repeats, eval):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
loss_args = EasyDict(
func_name='training.loss.ns_r1_DiffAugment') # Options for loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='4k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
metrics = [metric_defaults[x] for x in metrics]
metric_args = EasyDict(num_repeats=num_repeats)
desc = 'DiffAugment-stylegan2' if DiffAugment else 'stylegan2'
dataset_args = EasyDict(tfrecord_dir=dataset, resolution=resolution)
desc += '-' + os.path.basename(dataset)
if resolution is not None:
desc += '-{}'.format(resolution)
if num_samples is not None:
dataset_args.num_samples = num_samples
desc += '-{}samples'.format(num_samples)
if batch_size is not None:
desc += '-batch{}'.format(batch_size)
else:
batch_size = 32
assert batch_size % num_gpus == 0
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
sched.minibatch_size_base = batch_size
sched.minibatch_gpu_base = batch_size // num_gpus
G.impl = D.impl = impl
if fmap_base is not None:
G.fmap_base = D.fmap_base = fmap_base
desc += '-fmap{}'.format(fmap_base)
if fmap_max is not None:
G.fmap_max = D.fmap_max = fmap_max
desc += '-fmax{}'.format(fmap_max)
if latent_size is not None:
G.latent_size = G.mapping_fmaps = G.dlatent_size = latent_size
desc += '-latent{}'.format(latent_size)
if gamma is not None:
loss_args.gamma = gamma
desc += '-gamma{}'.format(gamma)
if DiffAugment:
loss_args.policy = DiffAugment
desc += '-' + DiffAugment.replace(',', '-')
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
loss_args=loss_args)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.update(resume_pkl=resume,
resume_kimg=resume_kimg,
resume_with_new_nets=True)
kwargs.update(metric_args=metric_args)
if ema_kimg is not None:
kwargs.update(G_ema_kimg=ema_kimg)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(**args):
args = EasyDict(args)
train = EasyDict(run_func_name = "training.training_loop.training_loop") # training loop options
sched = EasyDict() # TrainingSchedule options
vis = EasyDict() # visualize.eval() options
grid = EasyDict(size = "1080p", layout = "random") # setup_snapshot_img_grid() options
sc = dnnlib.SubmitConfig() # dnnlib.submit_run() options
# Environment configuration
tf_config = {
"rnd.np_random_seed": 1000,
"allow_soft_placement": True,
"gpu_options.per_process_gpu_memory_fraction": 1.0
}
if args.gpus != "":
num_gpus = len(args.gpus.split(","))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
# Networks configuration
cG = set_net("G", reg_interval = 4)
cD = set_net("D", reg_interval = 16)
# Dataset configuration
ratios = {
"clevr": 0.75,
"lsun-bedrooms": 0.72,
"cityscapes": 0.5,
"ffhq": 1.0
}
args.ratio = ratios.get(args.dataset, args.ratio)
dataset_args = EasyDict(tfrecord_dir = args.dataset, max_imgs = args.max_images, ratio = args.ratio,
num_threads = args.num_threads)
for arg in ["data_dir", "mirror_augment", "total_kimg"]:
cset(train, arg, args[arg])
# Training and Optimizations configuration
for arg in ["eval", "train", "recompile", "last_snapshots"]:
cset(train, arg, args[arg])
# Round to the closest multiply of minibatch size for validity
args.batch_size -= args.batch_size % args.minibatch_size
args.minibatch_std_size -= args.minibatch_std_size % args.minibatch_size
args.latent_size -= args.latent_size % args.component_num
if args.latent_size == 0:
print(bcolored("Error: latent-size is too small. Must best a multiply of component-num.", "red"))
exit()
sched_args = {
"G_lrate": "g_lr",
"D_lrate": "d_lr",
"minibatch_size": "batch_size",
"minibatch_gpu": "minibatch_size"
}
for arg, cmd_arg in sched_args.items():
cset(sched, arg, args[cmd_arg])
cset(train, "clip", args.clip)
# Logging and metrics configuration
metrics = [metric_defaults[x] for x in args.metrics]
cset(cG.args, "truncation_psi", args.truncation_psi)
for arg in ["summarize", "keep_samples"]:
cset(train, arg, args[arg])
# Visualization
args.imgs = args.images
args.ltnts = args.latents
vis_types ["imgs", "ltnts", "maps", "layer_maps", "interpolations", "noise_var", "style_mix"]:
# Set of all the set visualization types option
vis.vis_types = {arg for arg in vis_types if args[arg]}
vis_args = {
"grid": "vis_grid" ,
"num": "vis_num" ,
"rich_num": "vis_rich_num",
"section_size": "vis_section_size",
"intrp_density": "intrpolation_density",
"intrp_per_component": "intrpolation_per_component",
"alpha": "blending_alpha"
}
for arg, cmd_arg in vis_args.items():
cset(vis, arg, args[cmd_arg])
# Networks architecture
cset(cG.args, "architecture", args.g_arch)
cset(cD.args, "architecture", args.d_arch)
cset([cG.args, cD.args], "resnet_mlp", args.resnet_mlp)
cset(cG.args, "tanh", args.tanh)
# Latent sizes
if args.component_num > 1
if not (args.attention or args.merge):
print(bcolored("Error: component-num > 1 but the model is not using components.", "red"))
print(bcolored("Either add --attention for GANsformer or --merge for k-GAN).", "red"))
exit()
args.latent_size = int(args.latent_size / args.component_num)
cD.args.latent_size = cG.args.latent_size = cG.args.dlatent_size = args.latent_size
cset([cG.args, cD.args, train, vis], "component_num", args.component_num)
# Mapping network
for arg in ["layersnum", "lrmul", "dim", "shared"]:
cset(cG.args, arg, args["mapping_{}".formt(arg)])
# StyleGAN settings
for arg in ["style", "latent_stem", "fused_modconv", "local_noise"]:
cset(cG.args, arg, args[arg])
cD.args.mbstd_group_size = args.minibatch_std_size
# GANsformer
cset([cG.args, train], "attention", args.transformer)
cset(cD.args, "attention", args.d_transformer)
cset([cG.args, cD.args], "num_heads", args.num_heads)
args.norm = args.normalize
for arg in ["norm", "integration", "ltnt_gate", "img_gate", "kmeans",
"kmeans_iters", "asgn_direct", "mapping_ltnt2ltnt"]:
cset(cG.args, arg, args[arg])
for arg in ["attention_inputs", "use_pos"]:
cset([cG.args, cD.args], arg, args[arg])
# Positional encoding
for arg in ["dim", "init", "directions_num"]:
field = "pos_{}".format(arg)
cset([cG.args, cD.args], field, args[field])
# k-GAN
for arg in ["layer", "type", "channelwise"]:
field = "merge_{}".format(arg)
cset(cG.args, field, args[field])
cset([cG.args, train], "merge", args.merge)
# Attention
for arg in ["start_res", "end_res", "ltnt2ltnt", "img2img", "local_attention"]:
cset(cG.args, arg, args["g_{}".format(arg)])
cset(cD.args, arg, args["d_{}".format(arg)])
cset(cG.args, "img2ltnt", args.g_img2ltnt)
cset(cD.args, "ltnt2img", args.d_ltnt2img)
# Mixing and dropout
for arg in ["style_mixing", "component_mixing", "component_dropout", "attention_dropout"]:
cset(cG.args, arg, args[arg])
# Loss and regularization
gloss_args = {
"loss_type": "g_loss",
"reg_weight": "g_reg_weight"
"pathreg": "pathreg",
}
dloss_args = {
"loss_type": "d_loss",
"reg_type": "d_reg",
"gamma": "gamma"
}
for arg, cmd_arg in gloss_args.items():
cset(cG.loss_args, arg, args[cmd_arg])
for arg, cmd_arg in dloss_args.items():
cset(cD.loss_args, arg, args[cmd_arg])
##### Experiments management:
# Whenever we start a new experiment we store its result in a directory named 'args.expname:000'.
# When we rerun a training or evaluation command it restores the model from that directory by default.
# If we wish to restart the model training, we can set --restart and then we will store data in a new
# directory: 'args.expname:001' after the first restart, then 'args.expname:002' after the second, etc.
# Find the latest directory that matches the experiment
exp_dir = sorted(glob.glob("{}/{}:*".format(args.result_dir, args.expname)))[-1]
run_id = int(exp_dir.split(":")[-1])
# If restart, then work over a new directory
if args.restart:
run_id += 1
run_name = "{}:{0:03d}".format(args.expname, run_id)
train.printname = "{} ".format(misc.bold(args.expname))
snapshot, kimg, resume = None, 0, False
pkls = sorted(glob.glob("{}/{}/network*.pkl".format(args.result_dir, run_name)))
# Load a particular snapshot is specified
if args.pretrained_pkl:
# Soft links support
snapshot = glob.glob(args.pretrained_pkl)[0]
if os.path.islink(snapshot):
snapshot = os.readlink(snapshot)
# Extract training step from the snapshot if specified
try:
kimg = int(snapshot.split("-")[-1].split(".")[0])
except:
pass
# Find latest snapshot in the directory
elif len(pkls) > 0:
snapshot = pkls[-1]
kimg = int(snapshot.split("-")[-1].split(".")[0])
resume = True
if snapshot:
print(misc.bcolored("Resuming {}, kimg {}".format(snapshot, kimg), "white"))
train.resume_pkl = snapshot
train.resume_kimg = kimg
else:
print("Start model training from scratch.", "white")
# Run environment configuration
sc.run_dir_root = args.result_dir
sc.run_desc = args.expname
sc.run_id = run_id
sc.run_name = run_name
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(cG = cG, cD = cD)
kwargs.update(dataset_args = dataset_args, vis_args = vis, sched_args = sched, grid_args = grid, metric_arg_list = metrics, tf_config = tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.resume = resume
# If reload new options from the command line, no need to load the original configuration file
kwargs.load_config = not args.reload
dnnlib.submit_run(**kwargs)
metric_defaults = EasyDict([(args.name, args) for args in [
EasyDict(name='fid200-rt-shoes',
func_name='metrics.frechet_inception_distance.FID',
num_images=200,
minibatch_per_gpu=1,
ref_train=True,
ref_samples=49825),
EasyDict(name='fid200-rt-handbags',
func_name='metrics.frechet_inception_distance.FID',
num_images=200,
minibatch_per_gpu=1,
ref_train=True,
ref_samples=138567),
EasyDict(name='fid5k',
func_name='metrics.frechet_inception_distance.FID',
num_images=5000,
minibatch_per_gpu=8),
EasyDict(name='fid10k',
func_name='metrics.frechet_inception_distance.FID',
num_images=10000,
minibatch_per_gpu=8),
EasyDict(name='fid10k-b1',
func_name='metrics.frechet_inception_distance.FID',
num_images=10000,
minibatch_per_gpu=1),
EasyDict(name='fid10k-h0',
func_name='metrics.frechet_inception_distance.FID',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.0, 0.2]),
EasyDict(name='fid10k-h1',
func_name='metrics.frechet_inception_distance.FID',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.2, 0.4]),
EasyDict(name='fid10k-h2',
func_name='metrics.frechet_inception_distance.FID',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.4, 0.6]),
EasyDict(name='fid10k-h3',
func_name='metrics.frechet_inception_distance.FID',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.6, 0.8]),
EasyDict(name='fid10k-h4',
func_name='metrics.frechet_inception_distance.FID',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.8, 1.0]),
EasyDict(name='fid36k5',
func_name='metrics.frechet_inception_distance.FID',
num_images=36500,
minibatch_per_gpu=8),
EasyDict(name='fid36k5-h0',
func_name='metrics.frechet_inception_distance.FID',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.0, 0.2]),
EasyDict(name='fid36k5-h1',
func_name='metrics.frechet_inception_distance.FID',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.2, 0.4]),
EasyDict(name='fid36k5-h2',
func_name='metrics.frechet_inception_distance.FID',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.4, 0.6]),
EasyDict(name='fid36k5-h3',
func_name='metrics.frechet_inception_distance.FID',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.6, 0.8]),
EasyDict(name='fid36k5-h4',
func_name='metrics.frechet_inception_distance.FID',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.8, 1.0]),
EasyDict(name='fid50k',
func_name='metrics.frechet_inception_distance.FID',
num_images=50000,
minibatch_per_gpu=8),
EasyDict(name='ids5k',
func_name='metrics.inception_discriminator_score.IDS',
num_images=5000,
minibatch_per_gpu=8),
EasyDict(name='ids10k',
func_name='metrics.inception_discriminative_score.IDS',
num_images=10000,
minibatch_per_gpu=8),
EasyDict(name='ids10k-b1',
func_name='metrics.inception_discriminative_score.IDS',
num_images=10000,
minibatch_per_gpu=1),
EasyDict(name='ids10k-h0',
func_name='metrics.inception_discriminative_score.IDS',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.0, 0.2]),
EasyDict(name='ids10k-h1',
func_name='metrics.inception_discriminative_score.IDS',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.2, 0.4]),
EasyDict(name='ids10k-h2',
func_name='metrics.inception_discriminative_score.IDS',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.4, 0.6]),
EasyDict(name='ids10k-h3',
func_name='metrics.inception_discriminative_score.IDS',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.6, 0.8]),
EasyDict(name='ids10k-h4',
func_name='metrics.inception_discriminative_score.IDS',
num_images=10000,
minibatch_per_gpu=8,
hole_range=[0.8, 1.0]),
EasyDict(name='ids36k5',
func_name='metrics.inception_discriminative_score.IDS',
num_images=36500,
minibatch_per_gpu=8),
EasyDict(name='ids36k5-h0',
func_name='metrics.inception_discriminative_score.IDS',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.0, 0.2]),
EasyDict(name='ids36k5-h1',
func_name='metrics.inception_discriminative_score.IDS',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.2, 0.4]),
EasyDict(name='ids36k5-h2',
func_name='metrics.inception_discriminative_score.IDS',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.4, 0.6]),
EasyDict(name='ids36k5-h3',
func_name='metrics.inception_discriminative_score.IDS',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.6, 0.8]),
EasyDict(name='ids36k5-h4',
func_name='metrics.inception_discriminative_score.IDS',
num_images=36500,
minibatch_per_gpu=8,
hole_range=[0.8, 1.0]),
EasyDict(name='ids50k',
func_name='metrics.inception_discriminative_score.IDS',
num_images=50000,
minibatch_per_gpu=8),
EasyDict(
name='lpips2k',
func_name='metrics.learned_perceptual_image_patch_similarity.LPIPS',
num_pairs=2000,
minibatch_per_gpu=8),
]])
sess = tf.get_default_session()
sess.list_devices()
cores = tflex.get_cores()
tflex.set_override_cores(cores)
#synthesis_func = 'G_synthesis_stylegan2'
#kwargs = {'resolution': 512}
#synthesis = tflib.Network('G_synthesis', func_name=globals()[synthesis_func], **kwargs)
#sess.reset(os.environ['TPU_NAME']) # don't do this, this breaks the session
train = EasyDict(run_func_name='training.training_loop.training_loop') # Options for training loop.
G_args = EasyDict(func_name='training.networks_stylegan2.G_main') # Options for generator network.
D_args = EasyDict(func_name='training.networks_stylegan2.D_stylegan2') # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg') # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1') # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(size='8k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000}
label_dtype = np.int64
sched.minibatch_gpu = 1
if 'G' not in globals():
with tflex.device('/gpu:0'):
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma,
mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan2.G_main'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_r1'
) # Options for discriminator loss.
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
size='8k', layout='random') # Options for setup_snapshot_image_grid().
sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run().
tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf().
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 2
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
desc += '-' + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id != 'config-f':
G.fmap_base = D.fmap_base = 8 << 10
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {
128: 0.0015,
256: 0.002,
512: 0.003,
1024: 0.003
}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G,
D_args=D,
G_opt_args=G_opt,
D_opt_args=D_opt,
G_loss_args=G_loss,
D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args,
sched_args=sched,
grid_args=grid,
metric_arg_list=metrics,
tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, metrics):
train = EasyDict(run_func_name='training.training_loop.training_loop_mirror_v6_remove_half_fl_fr.training_loop')
G = EasyDict(func_name='training.networks.networks_stylegan2.G_main')
D = EasyDict(func_name='training.networks.networks_stylegan2_discriminator_new_rotation.D_stylegan2_new_rotaion') # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)
G_loss = EasyDict(func_name='training.loss.loss_G_new_rotation_squared_euclidean_10_interpolate_50_percent_uniform_dist_int_penalty.G_logistic_ns_pathreg')
D_loss = EasyDict(func_name='training.loss.loss_D_logistic_r1_new_rotation_euclidean_square.D_logistic_r1_new_rotation')
sched = EasyDict()
grid = EasyDict(size='1080p', layout='random')
sc = dnnlib.SubmitConfig()
tf_config = {'rnd.np_random_seed': 1000}
train.data_dir = data_dir
train.total_kimg = total_kimg
train.mirror_augment = mirror_augment
train.image_snapshot_ticks = train.network_snapshot_ticks = 10
sched.G_lrate_base = sched.D_lrate_base = 0.002
sched.minibatch_size_base = 32
sched.minibatch_gpu_base = 4
# train.resume_pkl = './results/00200-stylegan2-car_labels_v7_oversample_filter-2gpu-config-f-squared_euclidean_10_interpolate_50_percent_int_reg-256/network-snapshot-000887.pkl'
# train.resume_kimg = 887.2
D_loss.gamma = 10
metrics = [metric_defaults[x] for x in metrics]
desc = 'stylegan2'
G.style_mixing_prob = None
desc += '-' + dataset
dataset_args = EasyDict(tfrecord_dir=dataset)
assert num_gpus in [1, 2, 4, 8]
sc.num_gpus = num_gpus
desc += '-%dgpu' % num_gpus
assert config_id in _valid_configs
desc += '-' + config_id
desc += '-squared_euclidean_10_interpolate_50_percent_int_reg_remove_half_fl_fr_no_noise_square'
desc += '-256'
# Configs A-E: Shrink networks to match original StyleGAN.
if config_id != 'config-f':
G.fmap_base = D.fmap_base = 8 << 10
# Config E: Set gamma to 100 and override G & D architecture.
if config_id.startswith('config-e'):
D_loss.gamma = 100
if 'Gorig' in config_id: G.architecture = 'orig'
if 'Gskip' in config_id: G.architecture = 'skip' # (default)
if 'Gresnet' in config_id: G.architecture = 'resnet'
if 'Dorig' in config_id: D.architecture = 'orig'
if 'Dskip' in config_id: D.architecture = 'skip'
if 'Dresnet' in config_id: D.architecture = 'resnet' # (default)
# Configs A-D: Enable progressive growing and switch to networks that support it.
if config_id in ['config-a', 'config-b', 'config-c', 'config-d']:
sched.lod_initial_resolution = 8
sched.G_lrate_base = sched.D_lrate_base = 0.001
sched.G_lrate_dict = sched.D_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
sched.minibatch_size_base = 32 # (default)
sched.minibatch_size_dict = {8: 256, 16: 128, 32: 64, 64: 32}
sched.minibatch_gpu_base = 4 # (default)
sched.minibatch_gpu_dict = {8: 32, 16: 16, 32: 8, 64: 4}
G.synthesis_func = 'G_synthesis_stylegan_revised'
D.func_name = 'training.networks_stylegan2.D_stylegan'
# Configs A-C: Disable path length regularization.
if config_id in ['config-a', 'config-b', 'config-c']:
G_loss = EasyDict(func_name='training.loss.G_logistic_ns')
# Configs A-B: Disable lazy regularization.
if config_id in ['config-a', 'config-b']:
train.lazy_regularization = False
# Config A: Switch to original StyleGAN networks.
if config_id == 'config-a':
G = EasyDict(func_name='training.networks_stylegan.G_style')
D = EasyDict(func_name='training.networks_stylegan.D_basic')
if gamma is not None:
D_loss.gamma = gamma
sc.submit_target = dnnlib.SubmitTarget.LOCAL
sc.local.do_not_copy_source_files = True
kwargs = EasyDict(train)
kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss)
kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config)
kwargs.submit_config = copy.deepcopy(sc)
kwargs.submit_config.run_dir_root = result_dir
kwargs.submit_config.run_desc = desc
dnnlib.submit_run(**kwargs)
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""Main entry point for training StyleGAN and ProGAN networks."""
import copy
import dnnlib
from dnnlib import EasyDict
import config
from metrics import metric_base
# ----------------------------------------------------------------------------
# Official training configs for StyleGAN, targeted mainly for FFHQ.
if 1:
desc = 'sgan' # Description string included in result subdir name.
train = EasyDict(run_func_name='training.training_loop.training_loop'
) # Options for training loop.
G = EasyDict(func_name='training.networks_stylegan.G_style'
) # Options for generator network.
D = EasyDict(func_name='training.networks_stylegan.D_basic'
) # Options for discriminator network.
G_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for generator optimizer.
D_opt = EasyDict(beta1=0.0, beta2=0.99,
epsilon=1e-8) # Options for discriminator optimizer.
G_loss = EasyDict(func_name='training.loss.G_logistic_nonsaturating'
) # Options for generator loss.
D_loss = EasyDict(func_name='training.loss.D_logistic_simplegp',
r1_gamma=10.0) # Options for discriminator loss.
dataset = EasyDict() # Options for load_dataset().
sched = EasyDict() # Options for TrainingSchedule.
grid = EasyDict(
def G_synthesis_vc_modular(
dlatents_withl_in, # Input: Disentangled latents (W) [minibatch, label_size+dlatent_size].
dlatent_size=7, # Disentangled latent (W) dimensionality. Including discrete info, rotation, scaling, xy shearing, and xy translation.
label_size=0, # Label dimensionality, 0 if no labels.
module_list=None, # A list containing module names, which represent semantic latents (exclude labels).
num_channels=1, # Number of output color channels.
resolution=64, # Output resolution.
fmap_base=16 <<
10, # Overall multiplier for the number of feature maps.
fmap_decay=1.0, # log2 feature map reduction when doubling the resolution.
fmap_min=1, # Minimum number of feature maps in any layer.
fmap_max=512, # Maximum number of feature maps in any layer.
nonlinearity='lrelu', # Activation function: 'relu', 'lrelu', etc.
dtype='float32', # Data type to use for activations and outputs.
resample_kernel=[
1, 3, 3, 1
], # Low-pass filter to apply when resampling activations. None = no filtering.
fused_modconv=True, # Implement modulated_conv2d_layer() as a single fused op?
use_noise=False, # If noise is used in this dataset.
randomize_noise=True, # True = randomize noise inputs every time (non-deterministic), False = read noise inputs from variables.
single_const=True, # If only use a single constant feature at the begining.
where_feat_map=15, # For F_loss, which layer of feat map to use.
**_kwargs): # Ignore unrecognized keyword args.
'''
Modularized variation-consistent network.
'''
resolution_log2 = int(np.log2(resolution)) # == 6 for resolution 64
assert resolution == 2**resolution_log2 and resolution >= 4
def nf(stage):
return np.clip(int(fmap_base / (2.0**(stage * fmap_decay))), fmap_min,
fmap_max)
num_layers = resolution_log2 * 2 - 2 # == 10 for resolution 64
act = nonlinearity
images_out = None
# Note that module_list may include modules not containing latents,
# e.g. Conv layers (size in this case means number of conv layers).
key_ls, size_ls, count_dlatent_size, n_content = split_module_names(
module_list)
print('In key_ls:', key_ls)
print('In size_ls:', size_ls)
print('In count_dlatent_size:', count_dlatent_size)
if label_size > 0:
key_ls.insert(0, 'Label')
size_ls.insert(0, label_size)
n_content += label_size
# module_dict = collections.OrderedDict(zip(key_ls, size_ls))
# Primary inputs.
assert dlatent_size == count_dlatent_size
dlatents_withl_in.set_shape([None, label_size + count_dlatent_size])
dlatents_withl_in = tf.cast(dlatents_withl_in, dtype)
# Early layers consists of 4x4 constant layer.
y = None
if single_const:
with tf.variable_scope('4x4'):
with tf.variable_scope('Const'):
x = tf.get_variable(
'const',
shape=[1, 8, 4, 4],
initializer=tf.initializers.random_normal())
x = tf.tile(tf.cast(x, dtype),
[tf.shape(dlatents_withl_in)[0], 1, 1, 1])
else:
with tf.variable_scope('4x4'):
with tf.variable_scope('Const'):
x = tf.get_variable(
'const',
shape=[n_content, 8, 4, 4],
initializer=tf.initializers.random_normal())
subkwargs = EasyDict()
subkwargs.update(dlatents_withl_in=dlatents_withl_in,
n_content=n_content,
act=act,
dtype=dtype,
resample_kernel=resample_kernel,
fused_modconv=fused_modconv,
use_noise=use_noise,
randomize_noise=randomize_noise)
# Build modules by module_dict.
start_idx = 0
# print('module_dict:', module_dict)
# for scope_idx, k in enumerate(module_dict):
for scope_idx, k in enumerate(key_ls):
if scope_idx == where_feat_map:
feat_map = x
if (k.startswith('Label')) or (k.startswith('D_global')):
# e.g. {'Label': 3}, {'D_global': 3}
x = build_D_layers(x,
name=k,
n_latents=size_ls[scope_idx],
start_idx=start_idx,
scope_idx=scope_idx,
single_const=single_const,
fmaps=nf(scope_idx),
**subkwargs)
start_idx += size_ls[scope_idx]
elif k.startswith('C_global'):
# e.g. {'C_global': 2}
x = build_C_global_layers(x,
name=k,
n_latents=size_ls[scope_idx],
start_idx=start_idx,
scope_idx=scope_idx,
**subkwargs)
start_idx += size_ls[scope_idx]
elif k.startswith('C_nocond_global'):
# e.g. {'C_nocond_global': 2}
x = build_C_global_nocond_layers(x,
name=k,
n_latents=size_ls[scope_idx],
start_idx=start_idx,
scope_idx=scope_idx,
fmaps=nf(scope_idx),
**subkwargs)
start_idx += size_ls[scope_idx]
elif k.startswith('C_local_heat'):
# e.g. {'C_local_heat': 4}
x = build_local_heat_layers(x,
name=k,
n_latents=size_ls[scope_idx],
start_idx=start_idx,
scope_idx=scope_idx,
fmaps=nf(scope_idx),
**subkwargs)
start_idx += size_ls[scope_idx]
elif k.startswith('C_local_hfeat'):
# e.g. {'C_local_hfeat_size': 4}
x = build_local_hfeat_layers(x,
name=k,
n_latents=size_ls[scope_idx],
start_idx=start_idx,
scope_idx=scope_idx,
fmaps=nf(scope_idx),
**subkwargs)
start_idx += size_ls[scope_idx]
elif k.startswith('Noise'):
# e.g. {'Noise': 1}
x = build_noise_layer(x,
name=k,
n_layers=size_ls[scope_idx],
scope_idx=scope_idx,
fmaps=nf(scope_idx),
**subkwargs)
elif k.startswith('Conv'):
# e.g. {'Conv-up': 2}, {'Conv-id': 1}
x = build_conv_layer(x,
name=k,
n_layers=size_ls[scope_idx],
scope_idx=scope_idx,
fmaps=nf(scope_idx),
**subkwargs)
else:
raise ValueError('Unsupported module type: ' + k)
y = torgb(x, y, num_channels=num_channels)
images_out = y
assert images_out.dtype == tf.as_dtype(dtype)
return tf.identity(images_out,
name='images_out'), tf.identity(feat_map,
name='feat_map')
