Package: gan-lab
GAN Lab1024x1024 images generated from a StyleGAN trained on FFHQ (not yet fully trained):GAN Lab currently supports:StyleGAN (https://arxiv.org/pdf/1812.04948.pdf)ProGAN (https://arxiv.org/pdf/1710.10196.pdf)ResNet GANsEach GAN model's default settings emulates its most recent official implementation, but at the same time this package features a configuration file (config.py) where the user can quickly tune an extensive list of hyperparameter settings to his/her choosing.Comes with additional features such as supervised learning capabilities, easy-to-use methods for saving/loading pretrained models, flexible learning rate scheduling (and re-scheduling) capabilities, etc.This package aims for an intuitive API without sacrificing any complexity anywhere.$ pip install gan-labThis will install all necessary dependencies for you and will enable the option to use the package like an API (see "Jupyter Notebook (or Custom Script) Usage" below).If you do not wish to use the package like an API (i.e. you just want to install dependencies and then just use the repo by means of runningtrain.py, like shown below in the "Basic Usage on Command-line" section), you can run '$ pip install -r requirements.txt' instead.Basic Usage on Command-lineClone this repo, then simply run the following to configure your model & dataset and train your chosen model:$ python config.py [model] [--optional_kwargs]
$ python data_config.py [dataset] [dataset_dir] [--optional_kwargs]
$ python train.pyThe model will be saved into the "./gan_lab/models" directory by default.If you would like to see a list of what each argument does, run '$ python config.py [model] -h' or '$ python data_config.py [dataset] [dataset_dir] -h' on the command-line.NOTE: Make sure that all images you would like to use in your model are located directly inside thedataset_dirparent directory before runningdata_config.py. Any images within subdirectories ofdataset_dir(except for the subdirectories named "train" or "valid" that get created when you rundata_config.py) will not be used when training your model.StyleGAN Example:A StyleGAN Generator that yields 128x128 images can be created by running the following 3 lines. Below is a snapshot of images as the StyleGAN progressively grows. Ofcourse, this is not the only configuration that works:$ python config.py stylegan --loss=nonsaturating --gradient_penalty=R1 --res_samples=128 --num_main_iters=1071000 --nimg_transition=630000 --batch_size=8 --enable_cudnn_autotuner --num_workers=12
$ python data_config.py FFHQ path/to/datasets/ffhq --enable_mirror_augmentation
$ python train.pyBy default, image grids like the ones above are saved periodically during training into the "./gan_lab/samples" directory every 1,000 iterations (seeconfig.py).ProGAN Example:A ProGAN Generator that yields 128x128 images like the ones below can be created by running the following 3 lines. Ofcourse, this is not the only configuration that works:$ python config.py progan --res_samples=128 --num_main_iters=1050000 --batch_size=8
$ python data_config.py CelebA-HQ path/to/datasets/celeba_hq --enable_mirror_augmentation
$ python train.pyBy default, image grids of generator output are saved periodically during training into the "./gan_lab/samples" directory every 1,000 iterations (seeconfig.py).ResNet GAN Example:A ResNet GAN Generator can be created by running the following 3 lines (for example):$ python config.py resnetgan --lr_base=.00015
$ python data_config.py LSUN-Bedrooms path/to/datasets/lsun_bedrooms
$ python train.py[SAMPLES FOR RESNET GAN COMING SOON]Jupyter Notebook (or Custom Script) UsageRunningtrain.pyis just the very basic usage. This package can be imported and utilized in a modular manner as well (like an API). For example, often it's helpful to experiment inside a Jupyter Notebook, like in the example workflow below.First, configure your GAN to your choosing on the command-line (like explained above under the "Basic Usage on Command-line" section):$ python config.py stylegan
$ python data_config.py FFHQ path/to/datasets/ffhqThen, write a custom script or Jupyter Notebook cells:fromgan_labimportget_current_configurationfromgan_lab.utils.data_utilsimportprepare_dataset,prepare_dataloaderfromgan_lab.stylegan.learnerimportStyleGANLearner# get most recent configurations:config=get_current_configuration('config')data_config=get_current_configuration('data_config')# get DataLoader(s)train_ds,valid_ds=prepare_dataset(data_config)train_dl,valid_dl,z_valid_dl=prepare_dataloader(config,data_config,train_ds,valid_ds)# instantiate StyleGANLearner and train:learner=StyleGANLearner(config)learner.train(train_dl,valid_dl,z_valid_dl)# train for config.num_main_iters iterationslearner.config.num_main_iters=300000# this is one example of changing your instantiated learner's configurationslearner.train(train_dl,valid_dl,z_valid_dl)# train for another 300000 iterations# save your trained model:learner.save_model('path/to/models/stylegan_model.tar')# later on, you can load this saved model by instantiating the same learner and then running load_model:# learner = StyleGANLearner( config )# learner.load_model( 'path/to/models/stylegan_model.tar' )Some Advantages of Jupyter Notebook (there are many more than this):You have the flexibility to think about what to do with your trained model after its trained rather than all at once, such as:whether you want to save/load your trained modelwhat learner.config parameters you want to change before training againYou can always stop the kernel during training, do something else, and then resume again and it will workNOTEthat by default, the--num_workersargument inconfig.pyis set to data-loading from just 1 subprocess; setting this to a larger number (that still falls within the constraints of your CPU(s)) will speed up training significantly. :slightly_smiling_face:TODO:Multi-GPU supportTensorBoard capabilitiesFID, IS, and MS-SSIM metrics calculationIncorporate Spectral NormalizationIncorporate Self-attentionIncorporate improvements from StyleGAN2 paper (https://arxiv.org/pdf/1912.04958.pdf)TorchScript capabilities