site/en/tutorials/generative/cyclegan.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v1CUZ0dkOo_F"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qmkj-80IHxnd"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_xnMOsbqHz61"
      },
      "source": [
        "# CycleGAN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ds4o1h4WHz9U"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/tutorials/generative/cyclegan\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/generative/cyclegan.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/tutorials/generative/cyclegan.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/tutorials/generative/cyclegan.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ITZuApL56Mny"
      },
      "source": [
        "This notebook demonstrates unpaired image to image translation using conditional GAN's, as described in [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593), also known as CycleGAN. The paper proposes a method that can capture the characteristics of one image domain and figure out how these characteristics could be translated into another image domain, all in the absence of any paired training examples. \n",
        "\n",
        "This notebook assumes you are familiar with Pix2Pix, which you can learn about in the [Pix2Pix tutorial](https://www.tensorflow.org/tutorials/generative/pix2pix). The code for CycleGAN is similar, the main difference is an additional loss function, and the use of unpaired training data.\n",
        "\n",
        "CycleGAN uses a cycle consistency loss to enable training without the need for paired data. In other words, it can translate from one domain to another without a one-to-one mapping between the source and target domain. \n",
        "\n",
        "This opens up the possibility to do a lot of interesting tasks like photo-enhancement, image colorization, style transfer, etc. All you need is the source and the target dataset (which is simply a directory of images).\n",
        "\n",
        "![Output Image 1](images/horse2zebra_1.png)\n",
        "![Output Image 2](images/horse2zebra_2.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1_Y75QXJS6h"
      },
      "source": [
        "## Set up the input pipeline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5fGHWOKPX4ta"
      },
      "source": [
        "Install the [tensorflow_examples](https://github.com/tensorflow/examples) package that enables importing of the generator and the discriminator."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bJ1ROiQxJ-vY"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://github.com/tensorflow/examples.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lhSsUx9Nyb3t"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YfIk2es3hJEd"
      },
      "outputs": [],
      "source": [
        "import tensorflow_datasets as tfds\n",
        "from tensorflow_examples.models.pix2pix import pix2pix\n",
        "\n",
        "import os\n",
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import clear_output\n",
        "\n",
        "AUTOTUNE = tf.data.AUTOTUNE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iYn4MdZnKCey"
      },
      "source": [
        "## Input Pipeline\n",
        "\n",
        "This tutorial trains a model to translate from images of horses, to images of zebras. You can find this dataset and similar ones [here](https://www.tensorflow.org/datasets/catalog/cycle_gan). \n",
        "\n",
        "As mentioned in the [paper](https://arxiv.org/abs/1703.10593), apply random jittering and mirroring to the training dataset. These are some of the image augmentation techniques that avoids overfitting.\n",
        "\n",
        "This is similar to what was done in [pix2pix](https://www.tensorflow.org/tutorials/generative/pix2pix#load_the_dataset)\n",
        "\n",
        "* In random jittering, the image is resized to `286 x 286` and then randomly cropped to `256 x 256`.\n",
        "* In random mirroring, the image is randomly flipped horizontally i.e., left to right."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iuGVPOo7Cce0"
      },
      "outputs": [],
      "source": [
        "dataset, metadata = tfds.load('cycle_gan/horse2zebra',\n",
        "                              with_info=True, as_supervised=True)\n",
        "\n",
        "train_horses, train_zebras = dataset['trainA'], dataset['trainB']\n",
        "test_horses, test_zebras = dataset['testA'], dataset['testB']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2CbTEt448b4R"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = 1000\n",
        "BATCH_SIZE = 1\n",
        "IMG_WIDTH = 256\n",
        "IMG_HEIGHT = 256"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yn3IwqhiIszt"
      },
      "outputs": [],
      "source": [
        "def random_crop(image):\n",
        "  cropped_image = tf.image.random_crop(\n",
        "      image, size=[IMG_HEIGHT, IMG_WIDTH, 3])\n",
        "\n",
        "  return cropped_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "muhR2cgbLKWW"
      },
      "outputs": [],
      "source": [
        "# normalizing the images to [-1, 1]\n",
        "def normalize(image):\n",
        "  image = tf.cast(image, tf.float32)\n",
        "  image = (image / 127.5) - 1\n",
        "  return image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fVQOjcPVLrUc"
      },
      "outputs": [],
      "source": [
        "def random_jitter(image):\n",
        "  # resizing to 286 x 286 x 3\n",
        "  image = tf.image.resize(image, [286, 286],\n",
        "                          method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
        "\n",
        "  # randomly cropping to 256 x 256 x 3\n",
        "  image = random_crop(image)\n",
        "\n",
        "  # random mirroring\n",
        "  image = tf.image.random_flip_left_right(image)\n",
        "\n",
        "  return image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tyaP4hLJ8b4W"
      },
      "outputs": [],
      "source": [
        "def preprocess_image_train(image, label):\n",
        "  image = random_jitter(image)\n",
        "  image = normalize(image)\n",
        "  return image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VB3Z6D_zKSru"
      },
      "outputs": [],
      "source": [
        "def preprocess_image_test(image, label):\n",
        "  image = normalize(image)\n",
        "  return image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RsajGXxd5JkZ"
      },
      "outputs": [],
      "source": [
        "train_horses = train_horses.cache().map(\n",
        "    preprocess_image_train, num_parallel_calls=AUTOTUNE).shuffle(\n",
        "    BUFFER_SIZE).batch(BATCH_SIZE)\n",
        "\n",
        "train_zebras = train_zebras.cache().map(\n",
        "    preprocess_image_train, num_parallel_calls=AUTOTUNE).shuffle(\n",
        "    BUFFER_SIZE).batch(BATCH_SIZE)\n",
        "\n",
        "test_horses = test_horses.map(\n",
        "    preprocess_image_test, num_parallel_calls=AUTOTUNE).cache().shuffle(\n",
        "    BUFFER_SIZE).batch(BATCH_SIZE)\n",
        "\n",
        "test_zebras = test_zebras.map(\n",
        "    preprocess_image_test, num_parallel_calls=AUTOTUNE).cache().shuffle(\n",
        "    BUFFER_SIZE).batch(BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e3MhJ3zVLPan"
      },
      "outputs": [],
      "source": [
        "sample_horse = next(iter(train_horses))\n",
        "sample_zebra = next(iter(train_zebras))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4pOYjMk_KfIB"
      },
      "outputs": [],
      "source": [
        "plt.subplot(121)\n",
        "plt.title('Horse')\n",
        "plt.imshow(sample_horse[0] * 0.5 + 0.5)\n",
        "\n",
        "plt.subplot(122)\n",
        "plt.title('Horse with random jitter')\n",
        "plt.imshow(random_jitter(sample_horse[0]) * 0.5 + 0.5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0KJyB9ENLb2y"
      },
      "outputs": [],
      "source": [
        "plt.subplot(121)\n",
        "plt.title('Zebra')\n",
        "plt.imshow(sample_zebra[0] * 0.5 + 0.5)\n",
        "\n",
        "plt.subplot(122)\n",
        "plt.title('Zebra with random jitter')\n",
        "plt.imshow(random_jitter(sample_zebra[0]) * 0.5 + 0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvX8sKsfMaio"
      },
      "source": [
        "## Import and reuse the Pix2Pix models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cGrL73uCd-_M"
      },
      "source": [
        "Import the generator and the discriminator used in [Pix2Pix](https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/pix2pix/pix2pix.py) via the installed [tensorflow_examples](https://github.com/tensorflow/examples) package.\n",
        "\n",
        "The model architecture used in this tutorial is very similar to what was used in [pix2pix](https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/pix2pix/pix2pix.py). Some of the differences are:\n",
        "\n",
        "* Cyclegan uses [instance normalization](https://arxiv.org/abs/1607.08022) instead of [batch normalization](https://arxiv.org/abs/1502.03167).\n",
        "* The [CycleGAN paper](https://arxiv.org/abs/1703.10593) uses a modified `resnet` based generator. This tutorial is using a modified `unet` generator for simplicity.\n",
        "\n",
        "There are 2 generators (G and F) and 2 discriminators (X and Y) being trained here. \n",
        "\n",
        "* Generator `G` learns to transform image `X` to image `Y`. $(G: X -> Y)$\n",
        "* Generator `F` learns to transform image `Y` to image `X`. $(F: Y -> X)$\n",
        "* Discriminator `D_X` learns to differentiate between image `X` and generated image `X` (`F(Y)`).\n",
        "* Discriminator `D_Y` learns to differentiate between image `Y` and generated image `Y` (`G(X)`).\n",
        "\n",
        "![Cyclegan model](images/cyclegan_model.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8ju9Wyw87MRW"
      },
      "outputs": [],
      "source": [
        "OUTPUT_CHANNELS = 3\n",
        "\n",
        "generator_g = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm')\n",
        "generator_f = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm')\n",
        "\n",
        "discriminator_x = pix2pix.discriminator(norm_type='instancenorm', target=False)\n",
        "discriminator_y = pix2pix.discriminator(norm_type='instancenorm', target=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wDaGZ3WpZUyw"
      },
      "outputs": [],
      "source": [
        "to_zebra = generator_g(sample_horse)\n",
        "to_horse = generator_f(sample_zebra)\n",
        "plt.figure(figsize=(8, 8))\n",
        "contrast = 8\n",
        "\n",
        "imgs = [sample_horse, to_zebra, sample_zebra, to_horse]\n",
        "title = ['Horse', 'To Zebra', 'Zebra', 'To Horse']\n",
        "\n",
        "for i in range(len(imgs)):\n",
        "  plt.subplot(2, 2, i+1)\n",
        "  plt.title(title[i])\n",
        "  if i % 2 == 0:\n",
        "    plt.imshow(imgs[i][0] * 0.5 + 0.5)\n",
        "  else:\n",
        "    plt.imshow(imgs[i][0] * 0.5 * contrast + 0.5)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O5MhJmxyZiy9"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 8))\n",
        "\n",
        "plt.subplot(121)\n",
        "plt.title('Is a real zebra?')\n",
        "plt.imshow(discriminator_y(sample_zebra)[0, ..., -1], cmap='RdBu_r')\n",
        "\n",
        "plt.subplot(122)\n",
        "plt.title('Is a real horse?')\n",
        "plt.imshow(discriminator_x(sample_horse)[0, ..., -1], cmap='RdBu_r')\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0FMYgY_mPfTi"
      },
      "source": [
        "## Loss functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JRqt02lupRn8"
      },
      "source": [
        "In CycleGAN, there is no paired data to train on, hence there is no guarantee that the input `x` and the target `y` pair are meaningful during training. Thus in order to enforce that the network learns the correct mapping, the authors propose the cycle consistency loss.\n",
        "\n",
        "The discriminator loss and the generator loss are similar to the ones used in [pix2pix](https://www.tensorflow.org/tutorials/generative/pix2pix#build_the_generator)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cyhxTuvJyIHV"
      },
      "outputs": [],
      "source": [
        "LAMBDA = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q1Xbz5OaLj5C"
      },
      "outputs": [],
      "source": [
        "loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wkMNfBWlT-PV"
      },
      "outputs": [],
      "source": [
        "def discriminator_loss(real, generated):\n",
        "  real_loss = loss_obj(tf.ones_like(real), real)\n",
        "\n",
        "  generated_loss = loss_obj(tf.zeros_like(generated), generated)\n",
        "\n",
        "  total_disc_loss = real_loss + generated_loss\n",
        "\n",
        "  return total_disc_loss * 0.5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "90BIcCKcDMxz"
      },
      "outputs": [],
      "source": [
        "def generator_loss(generated):\n",
        "  return loss_obj(tf.ones_like(generated), generated)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5iIWQzVF7f9e"
      },
      "source": [
        "Cycle consistency means the result should be close to the original input. For example, if one translates a sentence from English to French, and then translates it back from French to English, then the resulting sentence should be the same as the  original sentence.\n",
        "\n",
        "In cycle consistency loss, \n",
        "\n",
        "* Image $X$ is passed via generator $G$ that yields generated image $\\hat{Y}$.\n",
        "* Generated image $\\hat{Y}$ is passed via generator $F$ that yields cycled image $\\hat{X}$.\n",
        "* Mean absolute error is calculated between $X$ and $\\hat{X}$.\n",
        "\n",
        "$$forward\\ cycle\\ consistency\\ loss: X -> G(X) -> F(G(X)) \\sim \\hat{X}$$\n",
        "\n",
        "$$backward\\ cycle\\ consistency\\ loss: Y -> F(Y) -> G(F(Y)) \\sim \\hat{Y}$$\n",
        "\n",
        "\n",
        "![Cycle loss](images/cycle_loss.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NMpVGj_sW6Vo"
      },
      "outputs": [],
      "source": [
        "def calc_cycle_loss(real_image, cycled_image):\n",
        "  loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n",
        "  \n",
        "  return LAMBDA * loss1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U-tJL-fX0Mq7"
      },
      "source": [
        "As shown above, generator $G$ is responsible for translating image $X$ to image $Y$. Identity loss says that, if you fed image $Y$ to generator $G$, it should yield the real image $Y$ or something close to image $Y$.\n",
        "\n",
        "If you run the zebra-to-horse model on a horse or the horse-to-zebra model on a zebra, it should not modify the image much since the image already contains the target class.\n",
        "\n",
        "$$Identity\\ loss = |G(Y) - Y| + |F(X) - X|$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "05ywEH680Aud"
      },
      "outputs": [],
      "source": [
        "def identity_loss(real_image, same_image):\n",
        "  loss = tf.reduce_mean(tf.abs(real_image - same_image))\n",
        "  return LAMBDA * 0.5 * loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G-vjRM7IffTT"
      },
      "source": [
        "Initialize the optimizers for all the generators and the discriminators."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iWCn_PVdEJZ7"
      },
      "outputs": [],
      "source": [
        "generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n",
        "generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n",
        "\n",
        "discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n",
        "discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aKUZnDiqQrAh"
      },
      "source": [
        "## Checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WJnftd5sQsv6"
      },
      "outputs": [],
      "source": [
        "checkpoint_path = \"./checkpoints/train\"\n",
        "\n",
        "ckpt = tf.train.Checkpoint(generator_g=generator_g,\n",
        "                           generator_f=generator_f,\n",
        "                           discriminator_x=discriminator_x,\n",
        "                           discriminator_y=discriminator_y,\n",
        "                           generator_g_optimizer=generator_g_optimizer,\n",
        "                           generator_f_optimizer=generator_f_optimizer,\n",
        "                           discriminator_x_optimizer=discriminator_x_optimizer,\n",
        "                           discriminator_y_optimizer=discriminator_y_optimizer)\n",
        "\n",
        "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
        "\n",
        "# if a checkpoint exists, restore the latest checkpoint.\n",
        "if ckpt_manager.latest_checkpoint:\n",
        "  ckpt.restore(ckpt_manager.latest_checkpoint)\n",
        "  print ('Latest checkpoint restored!!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rw1fkAczTQYh"
      },
      "source": [
        "## Training\n",
        "\n",
        "Note: This example model is trained for fewer epochs (10) than the paper (200) to keep training time reasonable for this tutorial. The generated images will have much lower quality."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NS2GWywBbAWo"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RmdVsmvhPxyy"
      },
      "outputs": [],
      "source": [
        "def generate_images(model, test_input):\n",
        "  prediction = model(test_input)\n",
        "    \n",
        "  plt.figure(figsize=(12, 12))\n",
        "\n",
        "  display_list = [test_input[0], prediction[0]]\n",
        "  title = ['Input Image', 'Predicted Image']\n",
        "\n",
        "  for i in range(2):\n",
        "    plt.subplot(1, 2, i+1)\n",
        "    plt.title(title[i])\n",
        "    # getting the pixel values between [0, 1] to plot it.\n",
        "    plt.imshow(display_list[i] * 0.5 + 0.5)\n",
        "    plt.axis('off')\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kE47ERn5fyLC"
      },
      "source": [
        "Even though the training loop looks complicated, it consists of four basic steps:\n",
        "\n",
        "* Get the predictions.\n",
        "* Calculate the loss.\n",
        "* Calculate the gradients using backpropagation.\n",
        "* Apply the gradients to the optimizer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KBKUV2sKXDbY"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_step(real_x, real_y):\n",
        "  # persistent is set to True because the tape is used more than\n",
        "  # once to calculate the gradients.\n",
        "  with tf.GradientTape(persistent=True) as tape:\n",
        "    # Generator G translates X -> Y\n",
        "    # Generator F translates Y -> X.\n",
        "    \n",
        "    fake_y = generator_g(real_x, training=True)\n",
        "    cycled_x = generator_f(fake_y, training=True)\n",
        "\n",
        "    fake_x = generator_f(real_y, training=True)\n",
        "    cycled_y = generator_g(fake_x, training=True)\n",
        "\n",
        "    # same_x and same_y are used for identity loss.\n",
        "    same_x = generator_f(real_x, training=True)\n",
        "    same_y = generator_g(real_y, training=True)\n",
        "\n",
        "    disc_real_x = discriminator_x(real_x, training=True)\n",
        "    disc_real_y = discriminator_y(real_y, training=True)\n",
        "\n",
        "    disc_fake_x = discriminator_x(fake_x, training=True)\n",
        "    disc_fake_y = discriminator_y(fake_y, training=True)\n",
        "\n",
        "    # calculate the loss\n",
        "    gen_g_loss = generator_loss(disc_fake_y)\n",
        "    gen_f_loss = generator_loss(disc_fake_x)\n",
        "    \n",
        "    total_cycle_loss = calc_cycle_loss(real_x, cycled_x) + calc_cycle_loss(real_y, cycled_y)\n",
        "    \n",
        "    # Total generator loss = adversarial loss + cycle loss\n",
        "    total_gen_g_loss = gen_g_loss + total_cycle_loss + identity_loss(real_y, same_y)\n",
        "    total_gen_f_loss = gen_f_loss + total_cycle_loss + identity_loss(real_x, same_x)\n",
        "\n",
        "    disc_x_loss = discriminator_loss(disc_real_x, disc_fake_x)\n",
        "    disc_y_loss = discriminator_loss(disc_real_y, disc_fake_y)\n",
        "  \n",
        "  # Calculate the gradients for generator and discriminator\n",
        "  generator_g_gradients = tape.gradient(total_gen_g_loss, \n",
        "                                        generator_g.trainable_variables)\n",
        "  generator_f_gradients = tape.gradient(total_gen_f_loss, \n",
        "                                        generator_f.trainable_variables)\n",
        "  \n",
        "  discriminator_x_gradients = tape.gradient(disc_x_loss, \n",
        "                                            discriminator_x.trainable_variables)\n",
        "  discriminator_y_gradients = tape.gradient(disc_y_loss, \n",
        "                                            discriminator_y.trainable_variables)\n",
        "  \n",
        "  # Apply the gradients to the optimizer\n",
        "  generator_g_optimizer.apply_gradients(zip(generator_g_gradients, \n",
        "                                            generator_g.trainable_variables))\n",
        "\n",
        "  generator_f_optimizer.apply_gradients(zip(generator_f_gradients, \n",
        "                                            generator_f.trainable_variables))\n",
        "  \n",
        "  discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients,\n",
        "                                                discriminator_x.trainable_variables))\n",
        "  \n",
        "  discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients,\n",
        "                                                discriminator_y.trainable_variables))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2M7LmLtGEMQJ"
      },
      "outputs": [],
      "source": [
        "for epoch in range(EPOCHS):\n",
        "  start = time.time()\n",
        "\n",
        "  n = 0\n",
        "  for image_x, image_y in tf.data.Dataset.zip((train_horses, train_zebras)):\n",
        "    train_step(image_x, image_y)\n",
        "    if n % 10 == 0:\n",
        "      print ('.', end='')\n",
        "    n += 1\n",
        "\n",
        "  clear_output(wait=True)\n",
        "  # Using a consistent image (sample_horse) so that the progress of the model\n",
        "  # is clearly visible.\n",
        "  generate_images(generator_g, sample_horse)\n",
        "\n",
        "  if (epoch + 1) % 5 == 0:\n",
        "    ckpt_save_path = ckpt_manager.save()\n",
        "    print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,\n",
        "                                                         ckpt_save_path))\n",
        "\n",
        "  print ('Time taken for epoch {} is {} sec\\n'.format(epoch + 1,\n",
        "                                                      time.time()-start))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1RGysMU_BZhx"
      },
      "source": [
        "## Generate using test dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KUgSnmy2nqSP"
      },
      "outputs": [],
      "source": [
        "# Run the trained model on the test dataset\n",
        "for inp in test_horses.take(5):\n",
        "  generate_images(generator_g, inp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ABGiHY6fE02b"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "This tutorial has shown how to implement CycleGAN starting from the generator and discriminator implemented in the [Pix2Pix](https://www.tensorflow.org/tutorials/generative/pix2pix) tutorial. As a next step, you could try using a different dataset from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/cycle_gan). \n",
        "\n",
        "You could also train for a larger number of epochs to improve the results, or you could implement the modified ResNet generator used in the [paper](https://arxiv.org/abs/1703.10593) instead of the U-Net generator used here."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "cyclegan.ipynb",
            "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}