Energy-based Conditional Generative Adversarial Network (ECGAN)

This is the code for the NeurIPS 2021 paper "A Unified View of cGANs with and without Classifiers". The repository is modified from StudioGAN. If you find our work useful, please consider citing the following paper:

@inproceedings{chen2021ECGAN,
  title   = {A Unified View of cGANs with and without Classifiers},
  author  = {Si-An Chen and Chun-Liang Li and Hsuan-Tien Lin},
  booktitle = {Advances in Neural Information Processing Systems},
  year    = {2021}
}

Please feel free to contact Si-An Chen if you have any questions about the code/paper.

Introduction

We propose a new Conditional Generative Adversarial Network (cGAN) framework called Energy-based Conditional Generative Adversarial Network (ECGAN) which provides a unified view of cGANs and achieves state-of-the-art results. We use the decomposition of the joint probability distribution to connect the goals of cGANs and classification as a unified framework. The framework, along with a classic energy model to parameterize distributions, justifies the use of classifiers for cGANs in a principled manner. It explains several popular cGAN variants, such as ACGAN, ProjGAN, and ContraGAN, as special cases with different levels of approximations. An illustration of the framework is shown below.

Requirements

• Anaconda
• Python >= 3.6
• 6.0.0 <= Pillow <= 7.0.0
• scipy == 1.1.0 (Recommended for fast loading of Inception Network)
• sklearn
• seaborn
• h5py
• tqdm
• torch >= 1.6.0 (Recommended for mixed precision training and knn analysis)
• torchvision >= 0.7.0
• tensorboard
• 5.4.0 <= gcc <= 7.4.0 (Recommended for proper use of adaptive discriminator augmentation module)

You can install the recommended environment as follows:

conda env create -f environment.yml -n studiogan

With docker, you can use:

docker pull mgkang/studiogan:0.1

Quick Start

• Train (-t) and evaluate (-e) the model defined in CONFIG_PATH using GPU 0

CUDA_VISIBLE_DEVICES=0 python3 src/main.py -t -e -c CONFIG_PATH

• Train (-t) and evaluate (-e) the model defined in CONFIG_PATH using GPUs (0, 1, 2, 3) and DataParallel

CUDA_VISIBLE_DEVICES=0,1,2,3 python3 src/main.py -t -e -c CONFIG_PATH

Try python3 src/main.py to see available options.

Dataset

• CIFAR10: StudioGAN will automatically download the dataset once you execute main.py.

• Tiny Imagenet, Imagenet, or a custom dataset:

  1. download Tiny Imagenet and Imagenet. Prepare your own dataset.
  2. make the folder structure of the dataset as follows:

┌── docs
├── src
└── data
    └── ILSVRC2012 or TINY_ILSVRC2012 or CUSTOM
        ├── train
        │   ├── cls0
        │   │   ├── train0.png
        │   │   ├── train1.png
        │   │   └── ...
        │   ├── cls1
        │   └── ...
        └── valid
            ├── cls0
            │   ├── valid0.png
            │   ├── valid1.png
            │   └── ...
            ├── cls1
            └── ...

Examples and Results

The src/configs directory contains config files used in our experiments.

CIFAR10 (3x32x32)

To train and evaluate ECGAN-UC on CIFAR10:

python3 src/main.py -t -e -c src/configs/CIFAR10/ecgan_v2_none_0_0p01.json

Tiny ImageNet (3x64x64)

To train and evaluate ECGAN-UC on Tiny ImageNet:

python3 src/main.py -t -e -c src/configs/TINY_ILSVRC2012/ecgan_v2_none_0_0p01.json --eval_type valid

ImageNet (3x128x128)

To train and evaluate ECGAN-UCE on ImageNet (~12 days on 8 NVIDIA V100 GPUs):

python3 src/main.py -t -e -l -sync_bn -c src/configs/ILSVRC2012/imagenet_ecgan_v2_contra_1_0p05.json --eval_type valid

Generated Images

Here are some selected images generated by ECGAN.