This implementation is intended to serve as an example of coding an image-to-image translation model in Flax. The codebase makes use of Gin for configuration and includes support for TensorBoard logging of both training metrics and images with Google's Common Loop Utils library.
CycleGAN was published in 2017 by Jun-Yan Zhu, Taesung Park, Phillip Isola and Alexei A. Efros. Since then, the model has been covered thoroughly in online tutorials and blog posts. You will have no trouble finding resources that explain how it works. Please start by looking at the original project page and reading the paper if you are not already familiar with it.
*Thanks to Barbara Olsen for providing the video that I used to make the animation to the right.
Weights for models trained on the horse2zebra and monet2photos datasets are published to HuggingFace at the following link: https://huggingface.co/louwjac/CycleGAN-Flax
The inference demo in the "notebooks" folder provides a detailed example of how to instantiate a model and load the trained weights from HuggingFace. It also includes a Gradio demo that makes it easy to submit external images to the model.
The TensorBoard logs for the training runs can be viewed at the links below. The logged images are unfortunately not available due to file size constraints.
This repo uses the TensorFlow Datasets version of the CycleGAN datasets. See ./cyclegan/data.py for the details on the input pipeline.
- Linux
- Local Python environment
- Nvidia GPU with the latest drivers
- Git
git clone https://github.com/louwjac/CycleGAN-Flax
cd CycleGAN-Flax
pip install -r requirements.txt
python main.py config/horse2zebra_original.ginGin is an extremely convenient configuration library that can replace other approaches to configuring parameters such as ConfigDict or Abseil flags. See the discussion on using Gin with Flax here: #1226, and keep it in mind when you use Gin for your own projects. Configuration files for several of the CycleGAN datasets are provided in the "config" folder of this repo.
Common Loop Utils (or CLU) is a Python library that provides tools to make it easier to write training loops. For the purposes of this project, that consists of utilities to capture and aggregate training metrics in addition to logging to both files and TensorBoard. CLU is a very young project and documentation is still scarce as of June, 2023. You'll need to read CLU's source code and look at code examples to learn how to use it. Please feel free to let me know if you have a question about the way that I used it.