Skip to content

TPU enabled Tensorflow implementation of "Large Scale GAN Training for High Fidelity Natural Image Synthesis" (BigGAN)

License

Notifications You must be signed in to change notification settings

Mistobaan/BigGAN-TPU-TensorFlow

 
 

Repository files navigation

BigGAN Tensorflow TPU

Simple Tensorflow TPU implementation of "Large Scale GAN Training for High Fidelity Natural Image Synthesis" (BigGAN)

I (David Mack) have been modifying this network to allow for configuration of its self-attention, to facilitate experiments into the effectiveness of different self-attention architectures.

main

Implementation notes/issues

  • TODO: Ensure BatchNorm is applied across all TPUs, not per-TPU
  • TODO: Implement BigGAN-deep architecture (simpler class embedding, deeper resblock)
  • TODO: Refactor BigGAN_256.py and BigGAN_512.py to TPU compatable code (for example, see BigGAN128.py)
  • TODO: Explore whether orthogonal initialization (paper's method) should be used instead of random normal initialization (current implementation)

Usage

Building the data

For ImageNet, use TensorFlow's build scripts to create TFRecord files of your chosen image size (e.g. 128x128). --tfr-format inception

You can also use the data build script from NVidia's Progressive Growing of GANs. --tfr-format progan

Training

You can train on a Google TPU by setting the name of your TPU as an env var and running one of the training scripts. For example,

  • TPU_NAME=node-1 pipenv run ./launch_train_tpu_b128.sh

You need to have your training data stored on a Google cloud bucket.

Architecture

128x128

256x256

512x512

Contributing

You're very welcome to! Submit a PR or contact the author(s)

Authors

Junho Kim, David Mack

About

TPU enabled Tensorflow implementation of "Large Scale GAN Training for High Fidelity Natural Image Synthesis" (BigGAN)

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 97.1%
  • Shell 1.8%
  • Jupyter Notebook 1.1%