bryanlimy / tf2-cyclegan

TensorFlow 2 implementation of CycleGAN with multi-GPU training.

CycleGAN with multi-GPUs training in TensorFlow 2

This repository provide a concise example on how to use tf.distribute.MirroredStrategy with custom training loops in TensorFlow 2. We adapt the CycleGAN (Zhu et. al. 2017) tutorials from Keras and TensorFlow and train the model with multiple GPUs. See towardsdatascience.com/multi-gpus-and-custom-training-loops-in-tensorflow-2-15b4b86b53bd for a detailed tutorial.

horse → zebra → horse	zebra → horse → zebra

1. Setup

• create virtual environment for the project

  conda create -n cyclegan python=3.8
  

• activate virtual environment

  conda activate cyclegan
  

• install required packages

  sh setup.sh
  

2. Run

• We use the horse2zebra dataset from TensorFlow Datasets by default.
• Training logs and checkpoints are stored in --output_dir
• We can use the following command to train the CycleGAN model on 2 GPUs and store the TensorBoard summary and checkpoints to runs/:

  CUDA_VISIBLE_DEVICES=0,1 python main.py --output_dir runs/ --epochs 200
  

• Use --help to see all available flags.

3. Result

• Use TensorBoard to inspect the training summary and plots

  tensorboard --logdir runs/cyclegan