class_08_CycleGAN_ex.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CycleGAN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the input pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/tensorflow/examples.git\n",
      "  Cloning https://github.com/tensorflow/examples.git to c:\\users\\public\\documents\\estsoft\\creatortemp\\pip-req-build-wx81flut\n",
      "  Resolved https://github.com/tensorflow/examples.git to commit e3849ccb6d981aad311f0174281c8575d5e21646\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n",
      "Requirement already satisfied: absl-py in c:\\users\\owner\\anaconda\\lib\\site-packages (from tensorflow-examples===e3849ccb6d981aad311f0174281c8575d5e21646-) (0.11.0)\n",
      "Requirement already satisfied: six in c:\\users\\owner\\anaconda\\lib\\site-packages (from tensorflow-examples===e3849ccb6d981aad311f0174281c8575d5e21646-) (1.15.0)\n",
      "Building wheels for collected packages: tensorflow-examples\n",
      "  Building wheel for tensorflow-examples (setup.py): started\n",
      "  Building wheel for tensorflow-examples (setup.py): finished with status 'done'\n",
      "  Created wheel for tensorflow-examples: filename=tensorflow_examples-e3849ccb6d981aad311f0174281c8575d5e21646_-py3-none-any.whl size=301874 sha256=20a2f94834f2e4852b29f48dc820dd9a8f000413a0830e5e83ea16785ebae94e\n",
      "  Stored in directory: C:\\Users\\Public\\Documents\\ESTsoft\\CreatorTemp\\pip-ephem-wheel-cache-_tz3yuok\\wheels\\eb\\19\\50\\2a4363c831fa12b400af86325a6f26ade5d2cdc5b406d552ca\n",
      "Failed to build tensorflow-examples\n",
      "Installing collected packages: tensorflow-examples\n",
      "  Running setup.py install for tensorflow-examples: started\n",
      "  Running setup.py install for tensorflow-examples: finished with status 'done'\n",
      "Successfully installed tensorflow-examples-e3849ccb6d981aad311f0174281c8575d5e21646-\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Ignoring invalid distribution -umpy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -illow (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -cipy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -umpy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -illow (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -cipy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/tensorflow/examples.git 'C:\\Users\\Public\\Documents\\ESTsoft\\CreatorTemp\\pip-req-build-wx81flut'\n",
      "  WARNING: Built wheel for tensorflow-examples is invalid: Metadata 1.2 mandates PEP 440 version, but 'e3849ccb6d981aad311f0174281c8575d5e21646-' is not\n",
      "WARNING: Ignoring invalid distribution -umpy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -illow (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -cipy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "  DEPRECATION: tensorflow-examples was installed using the legacy 'setup.py install' method, because a wheel could not be built for it. pip 23.1 will enforce this behaviour change. A possible replacement is to fix the wheel build issue reported above. Discussion can be found at https://github.com/pypa/pip/issues/8368\n",
      "WARNING: Ignoring invalid distribution -umpy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -illow (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -cipy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -umpy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -illow (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -cipy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -umpy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -illow (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -cipy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -umpy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -illow (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -cipy (c:\\users\\owner\\anaconda\\lib\\site-packages)\n",
      "\n",
      "[notice] A new release of pip is available: 23.0.1 -> 23.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    }
   ],
   "source": [
    "!pip install git+https://github.com/tensorflow/examples.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The history saving thread hit an unexpected error (OperationalError('database or disk is full',)).History will not be written to the database.\n"
     ]
    },
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'tensorflow_datasets'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-4-2850369e5165>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow_datasets\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtfds\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mtensorflow_examples\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmodels\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mpix2pix\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mpix2pix\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mos\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'tensorflow_datasets'"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Input Pipeline  \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.\n",
    "\n",
    ">As mentioned in the paper, 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\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": {},
   "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": {},
   "outputs": [],
   "source": [
    "BUFFER_SIZE = 1000\n",
    "BATCH_SIZE = 1\n",
    "IMG_WIDTH = 256\n",
    "IMG_HEIGHT = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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\n",
    "\n",
    "# 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\n",
    "\n",
    "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\n",
    "\n",
    "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": {},
   "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": {},
   "outputs": [],
   "source": [
    "sample_horse = next(iter(train_horses))\n",
    "sample_zebra = next(iter(train_zebras))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Import and reuse the Pix2Pix models\n",
    "Import the generator and the discriminator used in Pix2Pix via the installed tensorflow_examples package.\n",
    "\n",
    "The model architecture used in this tutorial is very similar to what was used in pix2pix. Some of the differences are:\n",
    "\n",
    "Cyclegan uses instance normalization instead of batch normalization.\n",
    "The CycleGAN paper uses a modified resnet based generator. This tutorial is using a modified unet generator for simplicity.\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. \n",
    "Generator F learns to transform image Y to image 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "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": {},
   "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": {},
   "source": [
    "### Loss functions  \n",
    ">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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LAMBDA = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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\n",
    "\n",
    "def generator_loss(generated):\n",
    "    return loss_obj(tf.ones_like(generated), generated)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 \n",
    " is passed via generator \n",
    " that yields generated image \n",
    ".\n",
    "Generated image \n",
    " is passed via generator \n",
    " that yields cycled image \n",
    ".\n",
    "Mean absolute error is calculated between \n",
    " and \n",
    ".\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "As shown above, generator \n",
    " is responsible for translating image \n",
    " to image \n",
    ". Identity loss says that, if you fed image \n",
    " to generator \n",
    ", it should yield the real image \n",
    " or something close to image \n",
    ".\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "source": [
    "Initialize the optimizers for all the generators and the discriminators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "source": [
    "### Checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "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": {},
   "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": {},
   "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": {},
   "source": [
    "## Generate using test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the trained model on the test dataset\n",
    "for inp in test_horses.take(5):\n",
    "  generate_images(generator_g, inp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6 (tensorflow)",
   "language": "python",
   "name": "jeff_heaton_t81_park"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}