PixelSNAIL

This is a Python3 / Tensorflow implementation of PixelSNAIL.

This code base is based on OpenAI's PixelCNN++ code.

Setup

To run this code you need the following:

• a machine with multiple GPUs
• Python3
• Numpy, TensorFlow

Training the model

Use the train.py script to train the model.

Pretrained model checkpoint

You can download our pretrained (TensorFlow) CIFAR10 model and ImageNet model

CIFAR10

python train.py \
       --data_set=cifar \
       --model=h12_pool2_smallkey \
       --nr_logistic_mix=10 \
       --nr_filters=256 \
       --batch_size=8 \
       --init_batch_size=8 \
       --dropout_p=0.5 \
       --polyak_decay=0.9995 \
       --save_interval=10

ImageNet

python train.py \
       --data_set=imagenet \
       --model=h12_noup_smallkey \
       --nr_logistic_mix=32 \
       --nr_filters=256 \
       --batch_size=8 \
       --init_batch_size=8 \
       --learning_rate=0.0001 \
       --dropout_p=0.0 \
       --polyak_decay=0.9997 \
       --save_interval=1