Train DCGAN with TPUs on Google Cloud
Screenshot below shows 64px images generated using the code provided.
This repo contains code to train an unconditional DCGAN (Radford et al 2017) using TPUs on Google Cloud. It is based on the DCGAN TPU example by the Google Tensorflow team with the following modifications
64*64
and 128*128
generation: Provide two model architectures (mainly additional layers) that support generating higher resolution images (64, 128).The convert_to_tfrecords
script accepts arguments for data directory (data_dir
) and output file (output_file
). Data directory is expected to have folders which contain images directly.
python convert_to_tfrecords --data_dir=images/cifar --output_file=images/cifar/train.tfrecords --image_size=128
Expected
images
├── cifar
├── train
└── train_image1.jpg
└── train_image2.jpg
└── test
└── test_image1.jpg
└── test_image2.jpg
git clone https://github.com/victordibia/tpuDCGAN
export GCS_BUCKET_NAME= <Your GCS Bucket>
python dcgan_main.py --tpu=$TPU_NAME --train_data_file=gs://$GCS_BUCKET_NAME/data/masks/train_masks.tfrecords --dataset=dcgan64 --train_steps=10000 --train_steps_per_eval=500 --model_dir=gs://$GCS_BUCKET_NAME/dcgan/masks/model --test_data_file=gs://$GCS_BUCKET_NAME/data/rand/test.tfrecords
Interested in generating masks? This repo contains two trained models (64px and 128px). You can use the generate script to generate images using any of the models. If you have your own trained DCGAN models (ckpt files) you can point the script to the model directory.
python generate_from_model.py --model_dir=models/masks/128/model.ckpt-15000 --image_size=128 --output_dir=models/masks/128 --random_seed=2