Skip to main content

StyleGAN, ProGAN, and ResNet GANs

Project description

GAN Zoo

Currently supports:

Each GAN model's default settings emulates its most recent official implementation, but at the same time this package features a simple interface (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, flexible learning rate scheduling (and re-scheduling) capabilities, an easy-to-use interface for saving/loading the model/learner state (e.g. for transfer learning or for evaluation at a later time), etc.

This package aims for an intuitive API without sacrificing any complexity anywhere.


In your virtual environment (e.g. a conda virtual environment), run:

$ pip install gan-zoo

This 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 running train.py), run '$ pip install -r requirements.txt' instead.

Basic Usage on Command-line

Clone 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.py

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 the dataset_dir parent directory before running data_config.py. Any images within subdirectories of dataset_dir (except for the subdirectories named "train" or "valid" that get created when you run data_config.py) will not be used when training your model.

ProGAN Example:

A ProGAN Generator that yields 128x128 images (higher resolutions coming once model is done training in Google Colab with 16 GB GPU Memory!) like the ones below can be created by running the following 3 lines:

$ 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.py


By default, the image grid above is saved periodically during training into the "./gan_zoo/samples" directory every 1000 iterations (see config.py).

StyleGAN Example:

A StyleGAN Generator can be created by running the following 3 lines (for example):

$ python config.py stylegan --batch_size=16
$ python data_config.py FFHQ path/to/datasets/ffhq --enable_mirror_augmentation
$ python train.py

[SAMPLES COMING SOON]

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 COMING SOON]

Jupyter Notebook (or Custom Script) Usage

Running train.py is 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 below.

First configure your GAN to your choosing on the command-line (like explained above under "Basic Usage on Command-line"):

$ python config.py stylegan
$ python data_config.py FFHQ path/to/datasets/ffhq

Then write a custom script or Jupyter Notebook cells:

from gan_zoo import get_current_configuration
from gan_zoo.utils.data_utils import prepare_dataset, prepare_dataloader
from gan_zoo.stylegan.learner import StyleGANLearner

# 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 iterations
learner.config.num_main_iters = 300000            # this is one example of changing your instantiated learner's configurations
learner.train( train_dl, valid_dl, z_valid_dl )   # train for another 300000 iterations

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 model/learner
    • what learner.config parameters you want to change before training again
  • You can always stop the kernel during training and then resume again and it will work

NOTE that by default, the --num_workers argument in config.py is 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 (will be implemented soon):

  • [ ] Multi-GPU support
  • [ ] TensorBoard capabilities
  • [ ] FID, IS, and MS-SSIM metrics calculation
  • [ ] Incorporate Spectral Normalization
  • [ ] Incorporate Self-attention

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for gan-zoo, version 0.1.3
Filename, size File type Python version Upload date Hashes
Filename, size gan-zoo-0.1.3.tar.gz (61.5 kB) File type Source Python version None Upload date Hashes View hashes

Supported by

Elastic Elastic Search Pingdom Pingdom Monitoring Google Google BigQuery Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page