GAN_distributed_training.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GAN_distributed_training.ipynb",
      "provenance": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyMqaJAn3C/47yJinA/9q8Ur",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/SaKi1309/ML_Notebooks/blob/main/GAN_distributed_training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this Notebook, TensorFlow's tutorial on the [DCGAN](https://www.tensorflow.org/tutorials/generative/dcgan) is adapted to be able to train the model on TPU. \n",
        "[Check this guide]() for details on distributed training on TPU. \n",
        "\n",
        "<big> **Don't forget to turn on TPU in Runtime->Change runtime type->TPU!** 💪</big>"
      ],
      "metadata": {
        "id": "Rkt9n7HMrYuz"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Notebook Settings\n",
        "The following Parameters can be adapted to influence the behavior of the notebook."
      ],
      "metadata": {
        "id": "3hjSmc-pflPm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "IMAGE_SHAPE = (28,28,1) #@param\n",
        "NOISE_DIM = 100 #@param {type:\"integer\"} \n",
        "GLOBAL_BATCH_SIZE = 256 #@param {type:\"integer\"}\n",
        "EPOCHS =  30 #@param {type:\"integer\"}\n",
        "NUM_EXAMPLES_TO_GENERATE = 16 #@param {type:\"integer\"}"
      ],
      "metadata": {
        "id": "0dcYZAFvffu8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"0\"></a>\n",
        "# 0 - Initial Steps"
      ],
      "metadata": {
        "id": "7BC7YU1WBrOY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install imageio\n",
        "!pip install git+https://github.com/tensorflow/docs\n",
        "!pip install DeepSaki"
      ],
      "metadata": {
        "id": "5F12Tca0B93h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T0ZreW7yBfqd"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import glob\n",
        "import imageio\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import os\n",
        "import PIL\n",
        "import time\n",
        "import random\n",
        "\n",
        "from IPython import display\n",
        "import DeepSaki\n",
        "\n",
        "tf.__version__\n",
        "#DeepSaki.utils.EnableXlaAcceleration()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "strategy, RUNTIME_ENVIRONMENT, hw_accelerator_handle = DeepSaki.utils.DetectHw()"
      ],
      "metadata": {
        "id": "YSqkHIpYj43I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "CONFIG = {\n",
        "    \"BUFFER_SIZE\":60000,\n",
        "    \"BATCH_SIZE_PER_REPLICA\":int(GLOBAL_BATCH_SIZE / strategy.num_replicas_in_sync),\n",
        "    \"GLOBAL_BATCH_SIZE\":GLOBAL_BATCH_SIZE,\n",
        "    \"REDUCTION_STRATEGY\":tf.keras.losses.Reduction.NONE,\n",
        "    \"EPOCHS\":EPOCHS,\n",
        "    \"NOISE_DIM\":NOISE_DIM,\n",
        "    \"NUM_EXAMPLES_TO_GENERATE\":NUM_EXAMPLES_TO_GENERATE,\n",
        "    \"NOISE_SHAPE\":(NOISE_DIM,),\n",
        "    #\"IMAGE_SHAPE\":(32,32,3)\n",
        "    \"IMAGE_SHAPE\": IMAGE_SHAPE\n",
        "}\n",
        "\n",
        "GLOBAL_PARAM = {\n",
        "    \"RANDOM_SEED\": 10\n",
        "}"
      ],
      "metadata": {
        "id": "eRzRuFZZXJAc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#Seed is used to have comparability between runs while maintaining random behavior!\n",
        "random.seed(GLOBAL_PARAM[\"RANDOM_SEED\"])\n",
        "np.random.seed(seed=GLOBAL_PARAM[\"RANDOM_SEED\"])\n",
        "\n",
        "# You will reuse this seed overtime (so it's easier)\n",
        "# to visualize progress in the animated GIF)\n",
        "latent_noise = tf.random.normal([CONFIG[\"NUM_EXAMPLES_TO_GENERATE\"], CONFIG[\"NOISE_DIM\"]])"
      ],
      "metadata": {
        "id": "HLHkT6j4CY_8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"1\"></a>\n",
        "# 1 - Dataset"
      ],
      "metadata": {
        "id": "XEk9hkCwCMzA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The dataset is loaded, pre-processed and finally, a distributed dataset is created."
      ],
      "metadata": {
        "id": "WHyIKkCfk4kQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
        "train_images = train_images.reshape(train_images.shape[0], *CONFIG[\"IMAGE_SHAPE\"]).astype('float32')\n",
        "test_images = test_images.reshape(test_images.shape[0], *CONFIG[\"IMAGE_SHAPE\"]).astype('float32')\n",
        "train_images = (train_images - 127.5) / 127.5  # Normalize the images to [-1, 1]\n",
        "test_images = (test_images - 127.5) / 127.5  # Normalize the images to [-1, 1]\n",
        "\n",
        "print(\"Shape of train_images: {}\".format(np.shape(train_images)))\n",
        "print(\"Shape of test_images: {}\".format(np.shape(test_images)))\n",
        "\n",
        "# Batch and shuffle the data\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(CONFIG[\"BUFFER_SIZE\"]).batch(CONFIG[\"GLOBAL_BATCH_SIZE\"],drop_remainder=True)\n",
        "test_dataset = tf.data.Dataset.from_tensor_slices(test_images).shuffle(CONFIG[\"BUFFER_SIZE\"]).batch(CONFIG[\"GLOBAL_BATCH_SIZE\"],drop_remainder=True)\n",
        "\n",
        "#distribute\n",
        "train_dataset_distributed = strategy.experimental_distribute_dataset(train_dataset)\n",
        "test_dataset_distributed = strategy.experimental_distribute_dataset(test_dataset)"
      ],
      "metadata": {
        "id": "jjZrSRF2CRQK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"2\"></a>\n",
        "# 2 - Model Definition"
      ],
      "metadata": {
        "id": "UEeS0e_ZCreu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"2-1\"></a>\n",
        "## 2.1. Generator"
      ],
      "metadata": {
        "id": "ZpWIJhjKC1Ll"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def Generator(inputShape = CONFIG[\"NOISE_SHAPE\"]):\n",
        "  # specify the input shape\n",
        "  input_tensor = tf.keras.layers.Input(shape = inputShape)\n",
        " \n",
        "  x = tf.keras.layers.Dense(7*7*256, use_bias=False)(input_tensor)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "  x = tf.keras.layers.LeakyReLU()(x)\n",
        "\n",
        "  x = tf.keras.layers.Reshape((7, 7, 256))(x)\n",
        "\n",
        "  x = tf.keras.layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False)(x)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "  x = tf.keras.layers.LeakyReLU()(x)\n",
        "\n",
        "  x = tf.keras.layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False)(x)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "  x = tf.keras.layers.LeakyReLU()(x)\n",
        "\n",
        "  x = tf.keras.layers.Conv2DTranspose(CONFIG[\"IMAGE_SHAPE\"][-1], (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')(x)\n",
        "\n",
        "  # create the model\n",
        "  model = tf.keras.Model(inputs=input_tensor, outputs=x, name = \"Generator\")\n",
        "\n",
        "  return model\n",
        "\n",
        "#Testcode\n",
        "model = Generator()\n",
        "model.summary()\n",
        "tf.keras.utils.plot_model(model, show_shapes=True, expand_nested=False, show_dtype=True, to_file=\"generator.png\")"
      ],
      "metadata": {
        "id": "ezAo4IGHKEi5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Important to note, that the reduction strategy is set to tf.keras.losses.Reduction.NONE. Additionall,in line 4, the loss is reduced by summing all values and is scaled by the global batch size, not the per_replica_batch_size."
      ],
      "metadata": {
        "id": "FqTHog4zlPla"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generator_loss(fake_output):\n",
        "  cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True,reduction=CONFIG[\"REDUCTION_STRATEGY\"])\n",
        "  loss = cross_entropy(tf.ones_like(fake_output), fake_output)\n",
        "  loss = tf.reduce_sum(loss) *(1. / CONFIG[\"GLOBAL_BATCH_SIZE\"])\n",
        "  return loss"
      ],
      "metadata": {
        "id": "i1R3vemhE51l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"2-2\"></a>\n",
        "## 2.2. Discriminator"
      ],
      "metadata": {
        "id": "sJ-2owhwC_iF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def Discriminator(inputShape = CONFIG[\"IMAGE_SHAPE\"]):\n",
        "  # specify the input shape\n",
        "  input_tensor = tf.keras.layers.Input(shape = inputShape)\n",
        "  x = tf.keras.layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same')(input_tensor)\n",
        "  x = tf.keras.layers.LeakyReLU()(x)\n",
        "  x = tf.keras.layers.Dropout(0.3)(x)\n",
        "\n",
        "  x = tf.keras.layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same')(x)\n",
        "  x = tf.keras.layers.LeakyReLU()(x)\n",
        "  x = tf.keras.layers.Dropout(0.3)(x)\n",
        "\n",
        "  x = tf.keras.layers.Flatten()(x)\n",
        "  x = tf.keras.layers.Dense(1)(x)\n",
        "\n",
        "  # create the model\n",
        "  model = tf.keras.Model(inputs=input_tensor, outputs=x, name = \"Discriminator\")\n",
        "\n",
        "  return model\n",
        "\n",
        "#Testcode\n",
        "model = Discriminator()\n",
        "model.summary()\n",
        "tf.keras.utils.plot_model(model, show_shapes=True, expand_nested=False, show_dtype=True, to_file=\"discriminator.png\")"
      ],
      "metadata": {
        "id": "FZnXDhBvI1PH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As for the generator loss, the discriminator loss sets the reduction strategy to NONE, reduces by summing all values and scales by the global batch size"
      ],
      "metadata": {
        "id": "c6HfP4A5l6Jo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def discriminator_loss(real_output, fake_output):\n",
        "  cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True,reduction=CONFIG[\"REDUCTION_STRATEGY\"])\n",
        "  real_loss = cross_entropy(tf.ones_like(real_output), real_output)\n",
        "  fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)\n",
        "  total_loss = real_loss + fake_loss\n",
        "  total_loss = tf.reduce_sum(total_loss) *(1. / CONFIG[\"GLOBAL_BATCH_SIZE\"])\n",
        "  return total_loss"
      ],
      "metadata": {
        "id": "2E8CfhmBFGK3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"2-3\"></a>\n",
        "## 2.3. GAN"
      ],
      "metadata": {
        "id": "FvND48VVPCSj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GAN(tf.keras.Model):\n",
        "  def __init__(self, generator, discriminator, strategy = None):\n",
        "    super(GAN, self).__init__()\n",
        "\n",
        "    self.generator = generator\n",
        "    self.discriminator = discriminator\n",
        "    self.strategy = strategy\n",
        "\n",
        "  def compile(self, generator_optimizer, discriminator_optimizer, generator_loss, discriminator_loss):\n",
        "    super(GAN, self).compile()\n",
        "    #super(VoloGAN, self).compile(steps_per_execution = 64)\n",
        "    self.generator_optimizer = generator_optimizer\n",
        "    self.discriminator_optimizer = discriminator_optimizer\n",
        "    self.generator_loss = generator_loss\n",
        "    self.discriminator_loss = discriminator_loss\n",
        "\n",
        "  def call(self, batch):\n",
        "    generated = self.generator(batch)\n",
        "    discriminated = self.discriminator(generated)\n",
        "    return tf.stack([generated, discriminated])\n",
        "\n",
        "  def print_summary(self):\n",
        "    def getNumberOfElements(listOfLayers):\n",
        "      numElements = 0\n",
        "      for layer in listOfLayers:\n",
        "        numElements += tf.reduce_prod(layer.shape)\n",
        "      return numElements\n",
        "\n",
        "    summaryString = \"\"\n",
        "    summaryString += \"--------------------------------------------------\\n\"\n",
        "    summaryString += \"------------------- Summary GAN ------------------\\n\"\n",
        "    summaryString += \"--------------------------------------------------\\n\"\n",
        "    \n",
        "    summaryString += \"\\n\"\n",
        "    summaryString += self.generator.name + \":\\n\"\n",
        "    gen_total = getNumberOfElements(self.generator.variables)\n",
        "    gen_trainable = getNumberOfElements(self.generator.trainable_variables)\n",
        "    gen_nonTrainable = gen_total - gen_trainable\n",
        "    summaryString += \"  Total Variables: {:,}\\n\".format(gen_total)\n",
        "    summaryString += \"  Trainable Variables: {:,}\\n\".format(gen_trainable)\n",
        "    summaryString += \"  Non-trainable Variables: {:,}\\n\".format(gen_nonTrainable)\n",
        "\n",
        "    summaryString += \"\\n\"\n",
        "    summaryString += self.discriminator.name + \":\\n\"\n",
        "    disc_total = getNumberOfElements(self.discriminator.variables)\n",
        "    disc_trainable = getNumberOfElements(self.discriminator.trainable_variables)\n",
        "    disc_nonTrainable = disc_total - disc_trainable\n",
        "    summaryString += \"  Total Variables: {:,}\\n\".format(disc_total)\n",
        "    summaryString += \"  Trainable Variables: {:,}\\n\".format(disc_trainable)\n",
        "    summaryString += \"  Non-trainable Variables: {:,}\\n\".format(disc_nonTrainable)\n",
        "\n",
        "    summaryString += \"_________________________________________________\\n\"\n",
        "\n",
        "    GAN_total = gen_total + disc_total\n",
        "    GAN_trainable = gen_trainable + disc_trainable\n",
        "    GAN_nonTrainable = gen_nonTrainable + disc_nonTrainable\n",
        "    summaryString += \"Total Variables: {:,}\\n\".format(GAN_total)\n",
        "    summaryString += \"Trainable Variables: {:,}\\n\".format(GAN_trainable)\n",
        "    summaryString += \"Non-trainable Variables: {:,}\\n\".format(GAN_nonTrainable)\n",
        "\n",
        "    summaryString += \"_________________________________________________\\n\"\n",
        "\n",
        "    print(summaryString)\n",
        "\n",
        "    return summaryString\n",
        "\n",
        "  def generate_and_save_images(self, epoch, test_input):\n",
        "    # Notice `training` is set to False.\n",
        "    # This is so all layers run in inference mode (batchnorm).\n",
        "    predictions = self.generator(test_input, training=False)\n",
        "\n",
        "    fig = plt.figure(figsize=(4, 4))\n",
        "    fig.suptitle(\"Epoch: {}\".format(epoch), fontsize=10)\n",
        "\n",
        "    for i in range(predictions.shape[0]):\n",
        "      plt.subplot(4, 4, i+1)\n",
        "      if CONFIG[\"IMAGE_SHAPE\"][-1]==1:\n",
        "        plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')\n",
        "      else:\n",
        "        plt.imshow(predictions[i, :, :, 0:3] * 0.5 + 0.5)\n",
        "      plt.axis('off')\n",
        "\n",
        "    with tf.device('/job:localhost'): \n",
        "      plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))\n",
        "    plt.show()\n",
        "\n",
        "  @tf.function\n",
        "  def train_step(self, batch):\n",
        "    noise = tf.random.normal([CONFIG[\"BATCH_SIZE_PER_REPLICA\"], CONFIG[\"NOISE_DIM\"]])\n",
        "\n",
        "    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
        "      generated_images = self.generator(noise, training=True)\n",
        "\n",
        "      real_output = self.discriminator(batch, training=True)\n",
        "      fake_output = self.discriminator(generated_images, training=True)\n",
        "\n",
        "      gen_loss = self.generator_loss(fake_output)\n",
        "      disc_loss = self.discriminator_loss(real_output, fake_output)\n",
        "\n",
        "    gradients_of_generator = gen_tape.gradient(gen_loss, self.generator.trainable_variables)\n",
        "    gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)\n",
        "\n",
        "    self.generator_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))\n",
        "    self.discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))\n",
        "    return [gen_loss, disc_loss]\n",
        "\n",
        "  @tf.function\n",
        "  def test_step(self, batch):\n",
        "    noise = tf.random.normal([CONFIG[\"BATCH_SIZE_PER_REPLICA\"], CONFIG[\"NOISE_DIM\"]])\n",
        "    generated_images = self.generator(noise, training=False)\n",
        "\n",
        "    real_output = self.discriminator(batch, training=False)\n",
        "    fake_output = self.discriminator(generated_images, training=False)\n",
        "\n",
        "    gen_loss = self.generator_loss(fake_output)\n",
        "    disc_loss = self.discriminator_loss(real_output, fake_output)\n",
        "\n",
        "    return [gen_loss, disc_loss]\n",
        "\n",
        "  @tf.function\n",
        "  def distributed_train_step(self, batch_train):\n",
        "    ''' \n",
        "    per_replica_loss_vector:  vector of shape [per_replica-loss_gen1, per_replica-loss_gen2, per_replica-loss_disc1, per_replica-loss_disc1]\n",
        "    reduced_loss_vector:      Vector of results of the different devices (for TPU 8)\n",
        "    \n",
        "    '''\n",
        "    per_replica_loss_vector = self.strategy.run(self.train_step, args=(batch_train,))\n",
        "\n",
        "    #reduce the result of the replicas for every loss value returned!\n",
        "    reduced_loss_vector = []\n",
        "    for per_replica_loss in per_replica_loss_vector:\n",
        "      reduced_loss_vector.append(self.strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None))\n",
        "    \n",
        "    return reduced_loss_vector\n",
        "\n",
        "  @tf.function\n",
        "  def distributed_test_step(self, batch_test):\n",
        "    ''' \n",
        "    per_replica_loss_vector:  vector of shape [per_replica-loss_gen1, per_replica-loss_gen2, per_replica-loss_disc1, per_replica-loss_disc1]\n",
        "    reduced_loss_vector:      Vector of results of the different devices (for TPU 8)\n",
        "    \n",
        "    '''\n",
        "    per_replica_loss_vector = self.strategy.run(self.test_step, args=(batch_test,))\n",
        "\n",
        "    #reduce the result of the replicas for every loss value returned!\n",
        "    reduced_loss_vector = []\n",
        "    for per_replica_loss in per_replica_loss_vector:\n",
        "      reduced_loss_vector.append(self.strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None))\n",
        "    \n",
        "    return reduced_loss_vector\n",
        "\n",
        "  def train(self, train_dataset_distributed, test_dataset_distributed, epochs):\n",
        "    start_training = time.time()\n",
        "    for epoch in range(epochs):\n",
        "      start_epoch = time.time()\n",
        "\n",
        "      # initialize training parameter\n",
        "      total_train_loss = [0., 0.]\n",
        "      num_train_batches = 0\n",
        "\n",
        "      for image_batch in train_dataset_distributed:\n",
        "        total_train_loss = tf.math.add(total_train_loss, self.distributed_train_step(image_batch))\n",
        "        num_train_batches += 1\n",
        "      total_train_loss = total_train_loss / float(num_train_batches)\n",
        "\n",
        "      if test_dataset_distributed is not None:\n",
        "        # initialize training parameter\n",
        "        total_test_loss = [0., 0.]\n",
        "        num_test_batches = 0\n",
        "\n",
        "        for image_batch in test_dataset_distributed:\n",
        "          total_test_loss = tf.math.add(total_test_loss, self.distributed_test_step(image_batch))\n",
        "          num_test_batches += 1\n",
        "        total_test_loss = total_test_loss / float(num_test_batches)\n",
        "\n",
        "      # Produce images for the GIF as you go\n",
        "      display.clear_output(wait=True)\n",
        "      self.generate_and_save_images(epoch + 1, latent_noise)\n",
        "\n",
        "      print ('Total training time is {0:.3f} sec'.format(time.time()-start_training))\n",
        "      print ('Time for epoch {0:5} is {1:.3f} sec'.format(epoch + 1, time.time()-start_epoch))\n",
        "      print(\"Train:\\nGen. Loss: {0:.6f} \\nDisc. Loss: {1:.6f}\".format(total_train_loss[0], total_train_loss[1]))\n",
        "      if test_dataset_distributed is not None:\n",
        "        print(\"Test:\\nGen. Loss: {0:.6f} \\nDisc. Loss: {1:.6f}\".format(total_test_loss[0], total_test_loss[1]))\n",
        "\n",
        "    # Generate after the final epoch\n",
        "    display.clear_output(wait=True)\n",
        "    self.generate_and_save_images(epochs, latent_noise)\n",
        "    print ('Total training time is {0:.3f} sec'.format(time.time()-start_training))\n"
      ],
      "metadata": {
        "id": "bb6OZjjIPFKG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"3\"></a>\n",
        "# 3 - Training"
      ],
      "metadata": {
        "id": "x7AYxGyvFWXV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "It'S important to note, that the model is initialized within the context of `strategy.scope()`"
      ],
      "metadata": {
        "id": "GSX7S4BwmxQV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with strategy.scope():\n",
        "  ganModel = GAN(\n",
        "      Generator(),\n",
        "      Discriminator(),\n",
        "      strategy\n",
        "      )\n",
        "ganModel.compile(\n",
        "    generator_optimizer = tf.keras.optimizers.Adam(1e-4),\n",
        "    discriminator_optimizer = tf.keras.optimizers.Adam(1e-4),\n",
        "    generator_loss = generator_loss, \n",
        "    discriminator_loss = discriminator_loss\n",
        ")\n",
        "ganModel.print_summary()\n",
        "\n"
      ],
      "metadata": {
        "id": "L236qzfxF9HN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ganModel.train(\n",
        "    train_dataset_distributed=train_dataset_distributed,\n",
        "    test_dataset_distributed = None,# test_dataset, \n",
        "    epochs = CONFIG[\"EPOCHS\"])"
      ],
      "metadata": {
        "id": "wYPrggrfiXKJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"4\"></a>\n",
        "# 4 - Visualization"
      ],
      "metadata": {
        "id": "6yj29EWSGDU1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Display a single image using the epoch number\n",
        "def display_image(epoch_no):\n",
        "  return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))"
      ],
      "metadata": {
        "id": "JHUfzTpVGC6z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "display_image(CONFIG[\"EPOCHS\"])"
      ],
      "metadata": {
        "id": "U0A9H4SPGPui"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "anim_file = 'dcgan.gif'\n",
        "\n",
        "with imageio.get_writer(anim_file, mode='I') as writer:\n",
        "  filenames = glob.glob('image*.png')\n",
        "  filenames = sorted(filenames)\n",
        "  for filename in filenames:\n",
        "    image = imageio.imread(filename)\n",
        "    writer.append_data(image)\n",
        "  image = imageio.imread(filename)\n",
        "  writer.append_data(image)"
      ],
      "metadata": {
        "id": "ow4OnzJWGUr1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow_docs.vis.embed as embed\n",
        "embed.embed_file(anim_file)"
      ],
      "metadata": {
        "id": "HdC24UlEGWLH"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}