Skip to content

louwjac/CycleGAN-Flax

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

CycleGAN in Jax + Flax

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.

Trained Models

Weights

Weights for models trained on the horse2zebra and monet2photos datasets are published to HuggingFace at the following link: https://huggingface.co/louwjac/CycleGAN-Flax

Colab Notebook

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.

Open In Colab

TensorBoard logs

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.

  1. horse2zebra
  2. monet2photo

Training Data

This repo uses the TensorFlow Datasets version of the CycleGAN datasets. See ./cyclegan/data.py for the details on the input pipeline.

Running locally

Prerequisites

  • Linux
  • Local Python environment
  • Nvidia GPU with the latest drivers
  • Git

Installation

  git clone https://github.com/louwjac/CycleGAN-Flax
  cd CycleGAN-Flax
  pip install -r requirements.txt

Train a model

python main.py config/horse2zebra_original.gin

Notes

Gin

Gin 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

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.

About

An implementation of CycleGAN in Jax and Flax

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published