Translate a Horse to a Zebra: CycleGAN
By: Hao He
Tags: Neural Networks, Towards Data Science, Convolutional Network, Machine Learning, Cyclegan

Translate a Horse to a Zebra: CycleGAN

Ideal horse-to-zebra translation achieved by cycleGAN. Image taken from the official cycleGAN Github page.

In this article, you will join our journey to build and train a cycleGAN(Cycle-Consistent Generative Adversarial Networks) model to turn horses into zebras. We will first take a look at the structure of a cycleGAN as well as the four loss functions implemented to train it. We will also have a glimpse at the microarchitecture of each discriminator and generator deployed in a cycleGAN. Lastly, we will talk about how to train a cycleGAN with some very impressive translation outcomes attached at the very end. Before you start reading, it would be better to have some basic understanding in convolutional neural networks and GAN. This article is a summary of the Advanced Predictive Modeling course project accomplished by Josh Artmann, Yiming Jin, Kevin Cheng and Hao He at UT Austin.

Project Overview

As it is named, the goal of our project is simple: to build a machine which takes a horse image and translate it to a zebra image. Ideally, our machine will only change the pattern on the skin of the horse, while leave other features — such as background, the shape of the horse, other objects in the image etc. — totally unchanged. This type of task is called Image-to-Image Translation, which is a pretty mature research topic under neural network where there already exists a lot of paper about it.

Traditionally, to train an Image-to-Image Translation model, we will need a dataset of “paired samples”, that is, if we want to transfer a painting to a real photo, we will need hundreds of paintings and the corresponding real-life photos which the artworks were painted from. However, a challenge we are facing is, it is impossible for us to find the “paired zebra images” for our horses since such things don’t exist at all. This is why we are using a model called cycleGAN which doesn’t require such “paired samples” in training to achieve the transformation. Introduced by Jun-yan Zhu and his colleagues at UC Berkeley in 2017, cycleGAN presents a new way to do Image-to-Image Translation. While it still requires a set of horse images and a set of zebra images in training, the horse and zebra images don’t need to directly correlated with each other anymore.

To train a horse-to-zebra cycleGAN model, we need two sets of data: real-life horse images and zebra images. In our project, we are using a well-organized dataset grabbed from the official Github page for cycleGAN. The dataset contains two training folders each has about 1,200 horse or zebra photos. These photos are used to train the model. In addition, another 10 horse photos and 10 zebra photos are save in the test folders of the dataset to evaluate model performance.

Introduction to CycleGAN

A traditional Image-to-Image Translation could be achieved through a single GAN model, which is formed by two neural network units: one called the Generator, the other one called the Discriminator. For a cycleGAN model, it is in fact constructed by two GAN models. Before you ask why we need two, let first talk about how these two GANs are configured inside a cycleGAN.

The reason it is called a cycleGAN is that we will use the two GAN models to construct a cycle for some purpose which we will talk about soon. To make such cycle, we need to make the two GAN models to do exactly the same work, that is to do image-to-image translation, but in reversed directions. Specifically, we want GAN A to transfer a horse image to a zebra image, while GAN B has the ability to transfer a zebra image to a horse image. Each of the two GAN models is formed by a Conditional Generator and a PatchGAN Discriminator. We will talk about the structures of the generator and discriminators in the following section. To sum up, in our tool box we have two GAN models. Each of them contains one generator and one discriminator.

Visualization of the structure of a cycleGAN. Please forgive my terrible handwritings. Image by author.

Wait a minute, why the hell do we need two?

Let’s go back to the question raised previously: if we are only doing one-way image translation, or, horse-to-zebra without zebra-to-horse translation, why do we need two GANs in our model while one of them can get the job done? Well, to answer this question, let’s go back to the major advantage of a cycleGAN: it won’t require pair samples anymore in the training process. However, having paired samples is not always a bad thing. One merit of paired samples is that they put strict restriction on the generated images so that except the desired adjustments (e.g. zebra pattern), the rest part of the generated images will be unchanged from the source images. Although we take this restriction for granted when we have paired samples because it comes naturally with the strongly correlated paired images, it is not the case anymore when we switch to unpaired samples where there is no correlation between source images and target images. It might cause trouble if we continue to train a single GAN on these unpaired samples, that is, at some point of the training we might realize that our GAN starts to generate random zebra images which have little or no correlation at all with the horse images we feed it with. This is for sure the situation we don’t want to encounter. Therefore, to get rid of it, the ingenious creators of cycleGAN invented a concept called Cycle Consistency which will put the restriction back on generated images through an additional loss function. The creators themselves have made an excellent explanation on cycle consistency in their paper:

We exploit the property that translation should be “cycle consistent”, in the sense that if we translate, e.g., a sentence from English to French, and then translate it back from French to English, we should arrive back at the original sentence. —Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks, 2017.

Say you want to translate a French sentence into English and you want to see how well you did the job. A good way to evaluate your translation is to re-translate the English sentence back to French and ideally, your outcome should be exactly the same with the original French sentence. The same logic could be applied here: to ensure our model is on the right way in keeping the generated image correlated with the input, we firstly let GAN A take a horse image and transfer it to a generated zebra image and then feed it directly to GAN B which , by its nature, will gives us back a generated horse image. Ideally, the final generated horse image should look exactly the same with the original horse image, in the same way the re-translated French sentence matches the original French sentence. If it turns out there is a significant difference between them, we then know GAN A might not did a good job and needs to be upgraded through training. All we need to do then is convert the difference between the two horse images to a loss function. This process is called Forward Cycle, a half of the cycle structure which focuses on evaluating GAN A. The other half of the cycle, which is called Backward Cycle, will have exactly the same workflow of its counterpart but in a reversed direction so that it could evaluates GAN B. Combining the two loss functions from Forward and Backward Cycles, we obtain a general loss function called Cycle Consistency Loss and we will involve it as a part of the final loss function which we will use to update the cycleGAN model.

Visualization of Cycle Consistency workflow. Image by author.

Identity Loss

With Cycle Consistency Loss implemented, it seems like we don’t need to worry about the naughtiness of our model anymore. However, from their experiments, the creators of CycleGAN found out that while the outcome zebras seem to have the same shape and other stuff with the input horses under cycle consistency restriction, the color profile consistency between inputs and outputs was not maintained very well by the model. To solve this problem, they invented another structure called Identity Mapping. This structure is again constructed by two components, just like Cycle Consistency, where one of them focuses on the performance of GAN A and the other one evaluates GAN B. To evaluate GAN A, for example, instead of feeding it with a horse image, we will give it a zebra image. By its nature, GAN A will generate a new zebra image based on the original zebra. Ideally, all the stuff laying on the two images, including the colors, should look exactly the same. Again, we convert the difference between the two images to a loss function and combine it with the loss function from the other components to form an Identity Loss function, which will also be a part of the final loss function.

Visualization of Identity Mapping workflow. Image by author.

Unit Architectures

After taking a look at the macrostructure of a cycleGAN, let’s dive into the microstructure of each discriminator and generator inside a cycleGAN.

Discriminator

Recall for a discriminator, its job is to take an image and output a binary outcome (True or False) which reflects its judgment on whether the input image is real or fake.

Structure of the discriminators deployed in cycleGAN is somewhat traditional: a fully convolutional neural networks with totally five layer blocks where each block contains a 2D Convolution Layer with kernel size of 4x4 and stride of 2, an Instance Normalization Layer and a Leaky ReLU Layer (except the output block which uses a Sigmoid Layer as activation). Every time an image is put into a discriminator, each of the first four layer blocks will shrinkage the size of the image into half and double the number of channels. Therefore after the fourth layer, the input will have a size of 16x16 with 512 channels (The model input has a size of 256x256). Lastly, the output layer block will collapse the 512 channels into one 16x16 matrix as the final output.

Structure of a PatchGAN Discriminator. Thanks to Sarah Wolf who created this visualization in her article.

Wait a minute, why are we having a 16x16 matrix instead of a binary True/False as the outcome? This is because in cycleGAN, we are using a structure called PatchGAN for the discriminator. Unlike a normal discriminator which takes an image as a whole to judge whether it is real or fake, a PatchGAN discriminator will first divide the input image into a number of “patches” and then make individual judgements on each of them. In our case, the input image is divided into 16x16 patches in the end where each patch has a size of 70x70. The elements in our output matrix stands for the possibility of real or fake prediction the model made on each patch. To obtain a final binary judgment, we could either simply average them or compare the output matrix with a matrix of expected values and calculate the difference. More details about PatchGAN could be found here.

Generator

For a generator, its mission is to generate a transformed image as output based on the input image. This is achieved by a neural network structure which is formed by three sections: an encoder, a transformer and a decoder.

The encoder shrinkages the size of input images while increase its number of channels. Similar to Discriminator, it is constructed by three layer blocks where each contains a 2D Convolution Layer, an Instance Normalization Layer and a Leaky ReLU Layer. The first layer block doesn’t change the size of image, it only creates 64 channels on the input. The next two layer blocks however, each shrinkages the input size by half while also double the number of channels. The adjusted input is then passed to the transformer.

The transformer adds desired features to the input while keep the size of it unchanged. It contains six residual net blocks, or ResNet for short. Each ResNet contains two layer blocks: the first layer block contains a 2D Convolution Layer (with stride=1), an Instance Normalization Layer and a Leaky ReLU Layer. The second layer block contains a 2D Convolution Layer (with stride=1) and an Instance Normalization Layer. The transformed input is then passed to the decoder.

The decoder enlarges the input back to the original size and collapses all channels into RGB to form the final output image. The enlarging process is achieved by a stack of two Transpose Convolution Layers. For simplicity, you could regard a transpose convolution layer as a combination of a 2D Upsampling layer and a 2D convolution layer with stride=1. Generally speaking, it will enlarge the size of input while decrease the number of channels. In the end, the 256x256 pixel data with 64 channels produced by the two transpose convolution layers will be fed to the output layer, which will collapse the channels into RGB and thus output it as the final output image.

Structure of a cycleGAN Generator. Thanks to Sarah Wolf who created this visualization in her article.

Train CycleGAN

The training process at each iteration for a cycleGAN model follows the general rule which is applied to other GAN models: train the discriminator alone, while train the generator using a composite model where the discriminator is involved but set untrainable. We will apply this rule twice to train our 2 GAN models.

Workflow of one training iteration for cycleGAN. Notice at each iteration, we update the generator first since we will use it to generate fake images to train the corresponding discriminator. Image by author.

We initially designed 100 epochs in the training process. Each epoch contains about 1,187 iterations, which is the size of our training dataset. At each iteration, one horse image and one zebra image is extracted from the dataset to train the model.

Update Weight for a Generator Through a Composite Model

The composite models we use here to train the generators in cycleGAN is more complicated than common cases. Traditionally, a composite model generates only one loss function — the adversarial loss which reflects how well the discriminator identifies fake images produced by the generator. For a cycleGAN generator, however, as we discussed before, there are three other loss functions should be involved — the identity loss and forward/backward consistency loss. As a result, in our composite model, there are four independent components which produces the four loss functions. While the adversarial loss is represented by L2 loss function, the rest three are represented by L1. Once the four loss function is calculated, they are weight averaged to obtain the final loss function where the cycle consistency loss has the highest weight while the effect from adversarial loss is downplayed. At the end of each iteration, we record the final loss function and use Adam SGD to update model weight of the generator.

Flowchart for the composite model used to train GAN A. Notice at every iteration the composite model is fed by two images: one horse image and one zebra image. Image by author.

Update Weight for a Discriminator Using Image Pool

To train a discriminator, at each iteration, we use one real image and one fake image generated by the corresponding generator. While we grab the real image from training dataset directly, we use an image pool to prepare the fake image.

An image pool is a set of at most 50 fake images produced by the generator in the past iterations. The goal of image pools is to prevent cycleGAN model from changing drastically from iteration to iteration, according to Sarah Wolf’s article. To achieve this, at each iteration, we will randomly pick one fake image from the pool to train the discriminator instead of feeding it with the one just made by the generator directly. A chart that shows how an image pool is implemented and updated is presented below.

Flowchart for an image pool used to select fake image for Discriminator A. Image by author.

Once we have the two images ready, all we need to do is put them to the discriminator separately and sum up the two loss functions represented by L2. We record the loss sum and update the weight of the discriminator using Adam SGD.

Workflow of one training iteration for Discriminator A. Image by author.

A nature of GAN models is that the values of loss functions for generator and discriminator will reach an equilibrium after significant amount of iterations, thus it is hard for us to tell whether the model has reached an optimal status only based on the two loss functions. To solve this problem, we can save the model and some generated images produced by the model after every five epochs. We then decide on when to stop training by reviewing the sample model outputs manually.

Sample outcomes for model trained after 30 epochs. Notice while the model was not sophisticated in zebrafication at this time, it started to figure out where it should put zebra patterns on. Image by author.

From our actual training process, we realized the model performance became stable after 50 epochs and we didn’t find out any observable improvement on transformed images after 70 epochs. We stopped the process at the 75th epoch and use the then updated model as our ultimate zebra-maker.

Screenshot of our training process at a very early stage. Image by author.

Outcome Display

The model obtained after 75 epochs in fact shows impressive ability to put zebra pattern to a horse naturally and precisely.

Image by author.

Personally, I’m mostly impressed by the first translation. You can see there is a cute lamb lays beneath the horse in the original image. Our model accurately recognized the shape of the horse and excluded the lamb and anything else in the background when trying to add zebra patterns on the walking horse. These four pairs of images demonstrate that our model is a good zebra maker when a complete shape of horse with an appropriate size is presented in the original photo.

Image by author.

Other than transforming horse to zebra, we are surprised that out model also has the power to “zebraficate” other creatures like humans.

Some Failures

However, from our test runs we also found that our model is not always powerful. It could get confused in some cases.

Cases where cycleGAN model failed to translate a horse to a zebra. Image by author.

As you might observed, our model performed not very well when the shape of the horse is not presented completely in an obvious way (at least from the perspective of a neural network model). The horses in the pre-transform photos here are either zoomed in or zoomed out to much so that our model failed to detect them.

Our model chose the Supreme Leader instead of his horse when doing translation. Image by author.

Translation failure also occurs when someone else is stacked too closed to the horse, such as the charming Mr. Kim Jong-un.

It turns out our model doesn’t watch cartoon very often. Image by author.

We also tried to feed cartoon horses to our model. However, it turns out our model was not very into cartoons. As shown above, it insisted to zebraficate Winnie’s red shirt instead of Eyeeyore, a depressed donkey who has, at least from our human perspective, a more horse-like shape.

Reference

Thanks for your reading! Please feel free to check out our Github page for the data and code behind this project.