bigsml / mxnet-gan

MXNet Implementation of DCGAN, Conditional GAN, pix2pix

pix2pix in MXNet

MXNet Implementation of various GAN, including: DCGAN [1], CGAN [2], Image-to-Image translation [3] (a.k.a. pix2pix)

The main focus of the repo is to implement a MXNet version of pix2pix for research purpose.
Please refer to This paper by Isola et al. for more detail.
Here is the original code implemented by Torch and PyTorch

This is a working repo initially served as the final project for UIUC ECE544NA.

Prerequisites

• Linux (Tested in Ubuntu 16.04)
• Python 2 (You may need to modify some codes if you are using Python 3)
• CPU or NVIDIA GPU + CUDA CuDNN

Getting Started

Installation

• Build MXNet from source (tested using MXNet version v.0.11.1).

  git clone --recursive https://github.com/apache/incubator-mxnet mxnet
  cd mxnet
  cp make/config.mk .
  vim config.mk  # You need to change configuration in order to enable cuda and cudnn 
  make -j8
  cd python
  sudo python setup.py install

• Clone this repo:

  git clone https://github.com/bowenc0221/MXNet-GAN
  cd MXNet-GAN

• Put MXNet python package into ./external/mxnet/$(MXNET_VERSION) and modify MXNET_VERSION in ./experiments/*.yaml to $(YOUR_MXNET_PACKAGE).
• Install python packages.

  pip install Cython
  pip install EasyDict
  pip install opencv-python

DCGAN train/test

• Train

  python dcgan/train.py --cfg experiments/dcgan/mnist_dcgan.yaml

• Test

  python dcgan/test.py --cfg experiments/dcgan/mnist_dcgan.yaml

• Warning
  • I only implemented dcgan for mnist. You may need to write your own data iterator for other dataset.
  • I did not tune parameter for dcgan. I only trained for 1 epoch!

CGAN train/test

• train

  python cgan/train.py --cfg experiments/cgan/mnist_cgan.yaml

• test

  python cgan/test.py --cfg experiments/cgan/mnist_cgan.yaml

• Warning
  • I only implemented dcgan for mnist. You may need to write your own data iterator for other dataset.
  • I did not tune parameter for dcgan. I only trained for 1 epoch!

pix2pix train/test

• Download a pix2pix dataset (e.g.facades):

  bash ./datasets/download_pix2pix_dataset.sh facades

  Please refer to pytorch-CycleGAN-and-pix2pix for dataset information.
• Train a model:
  • AtoB

    python pix2pix/train.py --cfg experiments/pix2pix/facades_pix2pix_AtoB.yaml

  • BtoA

    python pix2pix/train.py --cfg experiments/pix2pix/facades_pix2pix_BtoA.yaml

• Test a model:
  • AtoB

    python pix2pix/test.py --cfg experiments/pix2pix/facades_pix2pix_AtoB.yaml

  • BtoA

    python pix2pix/test.py --cfg experiments/pix2pix/facades_pix2pix_BtoA.yaml

• PatchGAN
  • You can use any PatchGAN listed in the paper by changing netD in configuration to 'n_layers' and set n_layers to any number from 0-6.
  • n_layers = 0: pixelGAN 1x1 discriminator
  • n_layers = 1: patchGAN 16x16 discriminator
  • n_layers = 3: patchGAN 70x70 discriminator (default setting in the paper)
  • n_layers = 6: imageGAN 256x256 discriminator
• Train pix2pix on your own dataset
  • I only implemented pix2pix for cityscapes and facades dataset but you can generalize easily to your own dataset.
  • Prepare pix2pix-datasets according to this link
  • Modify num_train and num_val in ./data/generate_train_val.py and run the script.
  • In configuration file, modify dataset part, as well as batchsize and number of epoch.

Results

facades

Ground Truth

AtoB

BtoA

cityscapes

Ground Truth

AtoB

BtoA

edges2shoes

Ground Truth

AtoB

BtoA

Citation

If you use this code for your research, here is a list of paper you can refer to:

@inproceedings{goodfellow2014generative,
  title={Generative adversarial nets},
  author={Goodfellow, Ian and Pouget-Abadie, Jean and Mirza, Mehdi and Xu, Bing and Warde-Farley, David and Ozair, Sherjil and Courville, Aaron and Bengio, Yoshua},
  booktitle={Advances in neural information processing systems},
  pages={2672--2680},
  year={2014}
}
@article{mirza2014conditional,
  title={Conditional generative adversarial nets},
  author={Mirza, Mehdi and Osindero, Simon},
  journal={arXiv preprint arXiv:1411.1784},
  year={2014}
}
@article{radford2015unsupervised,
  title={Unsupervised representation learning with deep convolutional generative adversarial networks},
  author={Radford, Alec and Metz, Luke and Chintala, Soumith},
  journal={arXiv preprint arXiv:1511.06434},
  year={2015}
}
@article{pix2pix2016,
  title={Image-to-Image Translation with Conditional Adversarial Networks},
  author={Isola, Phillip and Zhu, Jun-Yan and Zhou, Tinghui and Efros, Alexei A},
  journal={arxiv},
  year={2016}
}

Reference

[1] DCGAN: Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks
[2] CGAN: Conditional Generative Adversarial Nets
[3] pix2pix: Image-to-Image Translation with Conditional Adversarial Networks

Acknowledgments

Code is inspired by:
[1] MXNet GAN Tutorial
[2] MXNet DCGAN Example
[3] A MXNet W-GAN Code
[4] pytorch-CycleGAN-and-pix2pix
[5] Gluon GAN Tutorials