diff --git a/data.sh b/data.sh new file mode 100644 index 00000000..9b061b0a --- /dev/null +++ b/data.sh @@ -0,0 +1,62 @@ +#----------------------------------SETTINGS------------------------------------# + +PROJECT_ID='for-nikita' #ID проекта +BUCKET_NAME='for-nikita' #Имя Bucket +LOAD='DATASETS' #DATASETS (загрузка датасета из GS) \ + #или IMAGES (создание датасета и его загрузка на GS). +PATH_TO_DATASETS='style-GAN2/datasets/256-gray' #Путь до датасета на GS +DATASET_FOLDER="img-gray-256" # Каталог для сохранения датасета на GS +RESOLUTION=256 #Разрешение img +MODE="gray" #gray or rgb +COUNT=5000 #Кол-во изображений для тестового датасета. + #Не передавать аргумент для использования всех изображений + +#------------------------------------------------------------------------------# + +gcloud auth login +gcloud config set project $PROJECT_ID + +cd ./stylegan2/ + +if [ $LOAD = "IMAGES" ]; then + gsutil cp gs://$BUCKET_NAME/images.zip ./images.zip; + unzip -q ./images.zip -d ./; +fi + +echo $0 + +full_path=$(realpath $0) +dir_path=$(dirname $full_path) +echo $dir_path + +PATH_TO_IMG="${dir_path}/images/" + +if [ $LOAD = "IMAGES" ]; then + echo Обработка изображений + mkdir -p ./images/custom/; + python3 preprocessing.py --resolution $RESOLUTION\ + --mode $MODE\ + --count $COUNT\ + --path $PATH_TO_IMG +fi + +if [ $LOAD = "IMAGES" ]; then + echo Создание датасета + python3 dataset_tool.py create_from_images ./datasets/custom ./images/custom + rm -rf ./images + echo Копирование датасета на GS + gsutil -m cp -r ./datasets/custom/*.tfrecords gs://for-nikita/style-GAN2/datasets/$DATASET_FOLDER/ +fi + +if [ $LOAD = "DATASETS" ]; then + echo Копирование датасета из GS + mkdir -p ./datasets/custom/; + gsutil -m cp -r gs://${BUCKET_NAME}/${PATH_TO_DATASETS}/*.tfrecords ./datasets/custom; +fi + + + + + + + diff --git a/dnnlib/tflib/_cudacache/fused_bias_act_a16e693e7645428cdfec300292c5d95a.so b/dnnlib/tflib/_cudacache/fused_bias_act_a16e693e7645428cdfec300292c5d95a.so new file mode 100644 index 00000000..535ef4c8 Binary files /dev/null and b/dnnlib/tflib/_cudacache/fused_bias_act_a16e693e7645428cdfec300292c5d95a.so differ diff --git a/dnnlib/tflib/_cudacache/upfirdn_2d_648bee25f470bc610f6903c48bb8c647.so b/dnnlib/tflib/_cudacache/upfirdn_2d_648bee25f470bc610f6903c48bb8c647.so new file mode 100644 index 00000000..5e4c044a Binary files /dev/null and b/dnnlib/tflib/_cudacache/upfirdn_2d_648bee25f470bc610f6903c48bb8c647.so differ diff --git a/preprocessing.py b/preprocessing.py new file mode 100644 index 00000000..668d2854 --- /dev/null +++ b/preprocessing.py @@ -0,0 +1,64 @@ +import cv2 +from PIL import Image +from tqdm import tqdm +import numpy as np +import os +import random +import argparse +import matplotlib.pyplot as plt + +parser = argparse.ArgumentParser(description='Train model') +parser.add_argument("--resolution", type=int) +parser.add_argument("--mode", type=str) +parser.add_argument("--count", type=int) +parser.add_argument("--path", type=str) +args = parser.parse_args() + + +def rgba2rgb(rgba, background=(255, 255, 255)): + row, col, ch = rgba.shape + if ch == 3: + return rgba + assert ch == 4, 'RGBA image has 4 channels.' + + rgb = np.zeros((row, col, 3), dtype='float32') + r, g, b, a = rgba[:, :, 0], rgba[:, :, 1], rgba[:, :, 2], rgba[:, :, 3] + + R, G, B = background + + rgb[:, :, 0] = r * a + (1.0 - a) * R + rgb[:, :, 1] = g * a + (1.0 - a) * G + rgb[:, :, 2] = b * a + (1.0 - a) * B + + return np.asarray(rgb, dtype='uint8') + + +def rgb2gray(rgb): + return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140]) + + +def preprocessing(resolution=256, mode="gray", count=None, path=None): + files = os.listdir(path) + if not count: + count = len(files) + for i in tqdm(range(0, count)): + if files[i][-3:] != 'png': + continue + try: + img = plt.imread(os.path.join(path, files[i])) + img = rgba2rgb(img) + if mode == "gray": + img = rgb2gray(img) + img = cv2.resize(img, (resolution, resolution)) + cv2.imwrite("{}custom/{}.jpeg".format(path,files[i][:-4]), img) + except: + pass + + +if __name__ == "__main__": + preprocessing( + resolution=args.resolution, + mode=args.mode, + count=args.count, + path=args.path + ) \ No newline at end of file diff --git a/run_generator.py b/run_generator.py index 339796c9..2418c832 100755 --- a/run_generator.py +++ b/run_generator.py @@ -33,7 +33,12 @@ def generate_images(network_pkl, seeds, truncation_psi): z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component] tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width] images = Gs.run(z, None, **Gs_kwargs) # [minibatch, height, width, channel] - PIL.Image.fromarray(images[0], 'RGB').save(dnnlib.make_run_dir_path('seed%04d.png' % seed)) + if images[0].shape[2] == 1: + H, W = images[0].shape[0], images[0].shape[1] + img = images[0].reshape(H, W) + else: + img = images[0] + PIL.Image.fromarray(img).save(dnnlib.make_run_dir_path('seed%04d.png' % seed)) #---------------------------------------------------------------------------- @@ -68,7 +73,12 @@ def style_mixing_example(network_pkl, row_seeds, col_seeds, truncation_psi, col_ print('Saving images...') for (row_seed, col_seed), image in image_dict.items(): - PIL.Image.fromarray(image, 'RGB').save(dnnlib.make_run_dir_path('%d-%d.png' % (row_seed, col_seed))) + if image.shape == 1: + H, W = image.shape[0], image.shape[1] + img = image.reshape(H, W) + else: + img = image + PIL.Image.fromarray(image).save(dnnlib.make_run_dir_path('%d-%d.png' % (row_seed, col_seed))) print('Saving image grid...') _N, _C, H, W = Gs.output_shape diff --git a/run_training.py b/run_training.py index bc4c0a2b..c4072867 100755 --- a/run_training.py +++ b/run_training.py @@ -3,19 +3,34 @@ # This work is made available under the Nvidia Source Code License-NC. # To view a copy of this license, visit # https://nvlabs.github.io/stylegan2/license.html - + import argparse import copy import os import sys - + import dnnlib from dnnlib import EasyDict - + from metrics.metric_defaults import metric_defaults - + #---------------------------------------------------------------------------- +parser = argparse.ArgumentParser(description='Train model') +parser.add_argument("--dataset", type=str) +parser.add_argument("--data_dir", type=str) +parser.add_argument("--num_gpus", type=int) +parser.add_argument("--total_kimg", type=int) +parser.add_argument("--mirror_augment", type=bool) +parser.add_argument("--minibatch_size_base", type=int) +parser.add_argument("--resolution", type=int) +parser.add_argument("--config_id", type=str) +parser.add_argument("--gamma") +parser.add_argument("--resume_pkl", type=str) +parser.add_argument("--result_dir", type=str) + +args = parser.parse_args() + _valid_configs = [ # Table 1 'config-a', # Baseline StyleGAN @@ -24,172 +39,115 @@ 'config-d', # + Path length regularization 'config-e', # + No growing, new G & D arch. 'config-f', # + Large networks (default) - + # Table 2 'config-e-Gorig-Dorig', 'config-e-Gorig-Dresnet', 'config-e-Gorig-Dskip', 'config-e-Gresnet-Dorig', 'config-e-Gresnet-Dresnet', 'config-e-Gresnet-Dskip', 'config-e-Gskip-Dorig', 'config-e-Gskip-Dresnet', 'config-e-Gskip-Dskip', ] - + +#---------------------------------------------------------------------------- + +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 = args.data_dir +train.total_kimg = args.total_kimg +train.mirror_augment = args.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 = args.minibatch_size_base +sched.minibatch_gpu_base = 4 +D_loss.gamma = 10 +metrics = [] +metrics = [metric_defaults[x] for x in metrics] +desc = 'stylegan2' + +desc += '-' + args.dataset +dataset_args = EasyDict(tfrecord_dir=args.dataset, resolution=args.resolution) + +assert args.num_gpus in [1, 2, 4, 8] +sc.num_gpus = args.num_gpus +desc += '-%dgpu' % args.num_gpus + +config_id = args.config_id + +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 = 8 # (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 args.gamma is not None: + D_loss.gamma = args.gamma + +sc.submit_target = dnnlib.SubmitTarget.LOCAL +sc.local.do_not_copy_source_files = True + +if args.resume_pkl == ' ': + resume_pkl = None +else: + resume_pkl = args.resume_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') # 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 +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_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config) + kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, + metric_arg_list=metrics, tf_config=tf_config,resume_pkl=resume_pkl) kwargs.submit_config = copy.deepcopy(sc) - kwargs.submit_config.run_dir_root = result_dir + kwargs.submit_config.run_dir_root = args.result_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs) - -#---------------------------------------------------------------------------- - -def _str_to_bool(v): - if isinstance(v, bool): - return v - if v.lower() in ('yes', 'true', 't', 'y', '1'): - return True - elif v.lower() in ('no', 'false', 'f', 'n', '0'): - return False - else: - raise argparse.ArgumentTypeError('Boolean value expected.') - -def _parse_comma_sep(s): - if s is None or s.lower() == 'none' or s == '': - return [] - return s.split(',') - + #---------------------------------------------------------------------------- - -_examples = '''examples: - - # Train StyleGAN2 using the FFHQ dataset - python %(prog)s --num-gpus=8 --data-dir=~/datasets --config=config-f --dataset=ffhq --mirror-augment=true - -valid configs: - - ''' + ', '.join(_valid_configs) + ''' - -valid metrics: - - ''' + ', '.join(sorted([x for x in metric_defaults.keys()])) + ''' - -''' - -def main(): - parser = argparse.ArgumentParser( - description='Train StyleGAN2.', - epilog=_examples, - formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR') - parser.add_argument('--data-dir', help='Dataset root directory', required=True) - parser.add_argument('--dataset', help='Training dataset', required=True) - parser.add_argument('--config', help='Training config (default: %(default)s)', default='config-f', required=True, dest='config_id', metavar='CONFIG') - parser.add_argument('--num-gpus', help='Number of GPUs (default: %(default)s)', default=1, type=int, metavar='N') - parser.add_argument('--total-kimg', help='Training length in thousands of images (default: %(default)s)', metavar='KIMG', default=25000, type=int) - parser.add_argument('--gamma', help='R1 regularization weight (default is config dependent)', default=None, type=float) - parser.add_argument('--mirror-augment', help='Mirror augment (default: %(default)s)', default=False, metavar='BOOL', type=_str_to_bool) - parser.add_argument('--metrics', help='Comma-separated list of metrics or "none" (default: %(default)s)', default='fid50k', type=_parse_comma_sep) - - args = parser.parse_args() - - if not os.path.exists(args.data_dir): - print ('Error: dataset root directory does not exist.') - sys.exit(1) - - if args.config_id not in _valid_configs: - print ('Error: --config value must be one of: ', ', '.join(_valid_configs)) - sys.exit(1) - - for metric in args.metrics: - if metric not in metric_defaults: - print ('Error: unknown metric \'%s\'' % metric) - sys.exit(1) - - run(**vars(args)) - -#---------------------------------------------------------------------------- - + if __name__ == "__main__": main() - -#---------------------------------------------------------------------------- - + +#---------------------------------------------------------------------------- \ No newline at end of file diff --git a/test_nvcc b/test_nvcc new file mode 100755 index 00000000..4abb9c84 Binary files /dev/null and b/test_nvcc differ diff --git a/train.sh b/train.sh new file mode 100644 index 00000000..f0d08704 --- /dev/null +++ b/train.sh @@ -0,0 +1,44 @@ +#----------------------------------SETTINGS------------------------------------# + +DATASET='custom' #Каталог с датасетом +DATA_DIR='./datasets' #Путь до каталога с датасетом +CONFIG_ID='config-f' #Конфигурация модели. По умолчанию config-e - большая сеть +NUM_GPUS=1 #Кол-во gpu +TOTAL_KIMG=100000 #Общая продолжительность обучения, измеряемая тысячами + #реальных изображений на один цикл. +MIRROR_AUGMENT=false #Зеркальная аугментация +MINIBATCH_SIZE=8 #Размер мини-пакета. По умолчанию 32 +RESOLUTION=256 #Разрешение изображений +RESULT_DIR='./saves/' #Каталог для сохранения результатов оубчения +PATH_TO_MODEL='for-nikita/style-GAN2/models/network-snapshot-015056.pkl' #Путь до модели на GS +PRETRAINED=true #Если TRUE, то модель продолжает обучаться +TRAIN=true #TRUE-обучение + +#------------------------------------------------------------------------------# + +IFS="/" read -a SPLIT_PATH <<< $PATH_TO_MODEL + +if [ $PRETRAINED = true ]; then +mkdir ./model/ + gsutil -m cp -r gs://$PATH_TO_MODEL ./model/ + RESUME_PKL="./model/${SPLIT_PATH[-1]}" +else + RESUME_PKL=' ' +fi + + + +if [ $TRAIN = true ]; then + cd ./stylegan2 + python3 run_training.py --dataset $DATASET\ + --data_dir $DATA_DIR\ + --config_id $CONFIG_ID\ + --num_gpus $NUM_GPUS\ + --total_kimg $TOTAL_KIMG\ + --mirror_augment $MIRROR_AUGMENT\ + --minibatch_size_base $MINIBATCH_SIZE\ + --resolution $RESOLUTION\ + --result_dir $RESULT_DIR\ + --resume_pkl $RESUME_PKL +fi + diff --git a/upload.sh b/upload.sh new file mode 100644 index 00000000..b51f49bf --- /dev/null +++ b/upload.sh @@ -0,0 +1,10 @@ +DATETIME=`date +"%Y-%m-%d-%H:%M:%S"` + +SAVES=./saves/ + +cd $SAVES +files=(*) + +MODEL_FOLDER=${files[-1]} + +gsutil -m cp -r $MODEL_FOLDER/ gs://for-nikita/style-GAN2/models/$MODEL_FOLDER'-'$DATETIME \ No newline at end of file