PixelCNN++
TensorFlow 2 distributed implementation of PixelCNN++ [1].
Setup
To install the required dependencies you can simply run:
pip install requirements.txt
You will need a machine that preferablly has multiple GPU's. Training PixelCNN++ on modest datasets like Cifar10 can take days or weeks with multiple GPU's to obtains results comparable to the original paper.
Training
To train on a single GPU or CPU:
python main.py --config experiments/mnist.gin single
To train with multiple GPUs:
python main.py --config experiments/mnist.gin multigpu
You can change Gin parameters on the command line by using the --binding flag. You can use this flag multiple times. For example,
python main.py --config experiments/mnist.gin --binding "train.batch_size=128" multigpu
Custom Datasets
You can easily use a custom dataset with this implementation.
You just need to define a function that returns a tuple (train, test) which returns a train and test Tensorflow tf.data.Dataset object. You can then bind train.dataset_fn to your new dataset function using Gin.
Note you can use a dataset with any number of channels without any further modification. If you use one channel make sure that the image shape is still (W, H, 1).
Omissions
-
Currently, evaluation doesn't employ Polyak averaging over previous weights. This is planned with the Tensorflow Addons
tfa.optimizers.MovingAverage. -
I don't currently use the "autoregressive channel" discussed in the original paper.
-
WeightNormalization doesn't use data dependent initialization. This is planned in the future with Tensorflow Addons.