DCGAN_on_cifar10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DCGANs on CIFAR-10 data\n",
    "You can run this notebook on google colab.\n",
    "\n",
    "<div align=\"center\"><a href=\"https://colab.research.google.com/github/vstark21/DCGANs/blob/master/DCGAN_on_cifar10.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\"></a></div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1bOZjprhbGP_"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import numpy as np\n",
    "from IPython import display\n",
    "import matplotlib.pyplot as plt\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating functions - \n",
    "* `generator_model` - returns a model to generate images.\n",
    "* `discriminator_model` - returns a model to discriminate original images and generated images.\n",
    "* `generator_loss` - loss to optimize generator model so that it produces images like original images.\n",
    "* `discriminator_loss` - loss to optimize discriminator model so that it discriminates original images and generated images correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jYzntNnFaJGU"
   },
   "outputs": [],
   "source": [
    "def generator_model():\n",
    "\n",
    "    model = keras.Sequential()\n",
    "\n",
    "    model.add(keras.layers.Dense(4*4*1024, input_shape=[100, ]))\n",
    "    model.add(keras.layers.BatchNormalization())\n",
    "    model.add(keras.layers.ReLU())\n",
    "\n",
    "    model.add(keras.layers.Reshape((4, 4, 1024)))\n",
    "\n",
    "    model.add(keras.layers.Conv2DTranspose(512, (5, 5), strides=(2, 2), padding=\"same\"))\n",
    "    model.add(keras.layers.BatchNormalization())\n",
    "    model.add(keras.layers.ReLU())\n",
    "\n",
    "    model.add(keras.layers.Conv2DTranspose(256, (5, 5), strides=(2, 2), padding=\"same\"))\n",
    "    model.add(keras.layers.BatchNormalization())\n",
    "    model.add(keras.layers.ReLU())\n",
    "\n",
    "    model.add(keras.layers.Conv2DTranspose(3, (5, 5), strides=(2, 2), padding=\"same\", activation=\"tanh\"))\n",
    "    \n",
    "    return model\n",
    "\n",
    "def discriminator_model():\n",
    "\n",
    "    model = keras.Sequential([\n",
    "                             keras.layers.Conv2D(64, (5, 5), strides=(2, 2), input_shape=(32, 32, 3), padding=\"same\"),\n",
    "                               keras.layers.Dropout(0.3),\n",
    "                             keras.layers.LeakyReLU(alpha=0.2),\n",
    "\n",
    "                             keras.layers.Conv2D(128, (5, 5), strides=(2, 2), padding=\"same\"),\n",
    "                              keras.layers.BatchNormalization(),\n",
    "                             keras.layers.Dropout(0.3),\n",
    "                             keras.layers.LeakyReLU(alpha=0.2),\n",
    "                              \n",
    "                             keras.layers.Conv2D(256, (5, 5), strides=(2, 2), padding=\"same\"),\n",
    "                              keras.layers.BatchNormalization(),\n",
    "                             keras.layers.Dropout(0.3),\n",
    "                             keras.layers.LeakyReLU(alpha=0.2),\n",
    "                              \n",
    "                             keras.layers.Conv2D(512, (5, 5), strides=(2, 2), padding=\"same\"),\n",
    "                             keras.layers.Dropout(0.4),\n",
    "                             keras.layers.LeakyReLU(alpha=0.2),\n",
    "\n",
    "                             keras.layers.Flatten(),\n",
    "                             keras.layers.Dense(1, activation=\"sigmoid\")\n",
    "    ])\n",
    "\n",
    "    return model\n",
    "\n",
    "def generator_loss(generated_output):\n",
    "\n",
    "    cross_entropy = keras.losses.BinaryCrossentropy()\n",
    "    loss = cross_entropy(tf.ones_like(generated_output), generated_output)\n",
    "    \n",
    "    return loss\n",
    "\n",
    "\n",
    "def discriminator_loss(original_output, generated_output):\n",
    "\n",
    "    cross_entropy = keras.losses.BinaryCrossentropy()\n",
    "    \n",
    "    real_loss =  cross_entropy(tf.ones_like(original_output), original_output)\n",
    "    gen_loss = cross_entropy(tf.zeros_like(generated_output), generated_output)\n",
    "\n",
    "    total_loss = real_loss + gen_loss\n",
    "\n",
    "    return total_loss\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading data from `tf.keras.datasets.cifar10` then reshaping, normalizing them and slicing them into minibatches.\n",
    "\n",
    "**NOTE :** \n",
    "> All the training images should be normalized to [-1, 1] beacuse all the discminator's inputs must have values which are in same range as generator model's output is computed using *tanh* so it returns values in range [-1, 1]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "36lyyDYhbH5Z"
   },
   "outputs": [],
   "source": [
    "(train_images, train_labels), (_, _) = tf.keras.datasets.cifar10.load_data()\n",
    "\n",
    "train_images = train_images.reshape(train_images.shape[0], 32, 32, 3).astype('float32')\n",
    "train_images = (train_images - 127.5) / 127.5 # Normalize the images to [-1, 1]\n",
    "\n",
    "BUFFER_SIZE = 60000\n",
    "BATCH_SIZE = 128\n",
    "\n",
    "# Batch and shuffle the data\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating a function `generate_and_save_images` so that we can save images after each epoch and creating a constant `TEST_INPUT` so that we can always use it to generate images to save."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FRDVXL1mbumu"
   },
   "outputs": [],
   "source": [
    "test_input = tf.random.normal((16, 100))\n",
    "def generate_and_save_images(model, epoch, test_input):\n",
    "\n",
    "    predictions = model(test_input, training=False) * 127.5 + 127.5\n",
    "    predictions = np.array(predictions, dtype=np.int32)\n",
    "\n",
    "    fig = plt.figure(figsize=(4,4))\n",
    "    for i in range(predictions.shape[0]):\n",
    "        plt.subplot(4, 4, i+1)\n",
    "        plt.imshow(predictions[i, :, :, :], cmap='gray')\n",
    "        plt.axis('off')\n",
    "\n",
    "    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2TINENU2b8cd"
   },
   "outputs": [],
   "source": [
    "EPOCHS = 100\n",
    "GEN_INPUT_DIMS = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating models and optimizers and training them using **mini-batch Gradient Descent**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hdOqRRwJb2vv"
   },
   "outputs": [],
   "source": [
    "generator = generator_model()\n",
    "discriminator = discriminator_model()\n",
    "generator_optimizer = keras.optimizers.Adam(0.0002, beta_1=0.5)\n",
    "discriminator_optimizer = keras.optimizers.Adam(0.0002, beta_1=0.5)\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "\n",
    "    print(f\"In epoch : {epoch}\")\n",
    "    epoch_start = time.time()\n",
    "\n",
    "    for batch in train_dataset:\n",
    "\n",
    "        noise = tf.random.normal([BATCH_SIZE, GEN_INPUT_DIMS])\n",
    "\n",
    "        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
    "\n",
    "            generated_images = generator(noise, training=True)\n",
    "            \n",
    "            original_output = discriminator(batch, training=True)\n",
    "            generated_output = discriminator(generated_images, training=True)\n",
    "\n",
    "            gen_loss = generator_loss(generated_output)\n",
    "            disc_loss = discriminator_loss(original_output, generated_output)\n",
    "\n",
    "        gen_gradients = gen_tape.gradient(gen_loss, generator.trainable_variables)\n",
    "        disc_gradients = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n",
    "\n",
    "        generator_optimizer.apply_gradients(zip(gen_gradients, generator.trainable_variables))\n",
    "        discriminator_optimizer.apply_gradients(zip(disc_gradients, discriminator.trainable_variables))\n",
    "        \n",
    "    display.clear_output(wait=True)\n",
    "    generate_and_save_images(generator, epoch, test_input)\n",
    "    print(f\"Epoch {epoch} took {time.time() - epoch_start} with gen_loss of {gen_loss} and disc_loss of {disc_loss}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating a *gif* using all the images generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rJVbuwv1foLl"
   },
   "outputs": [],
   "source": [
    "import os, imageio\n",
    "\n",
    "files = os.listdir()\n",
    "\n",
    "images = []\n",
    "\n",
    "for filename in files:\n",
    "\n",
    "    if filename[-3:] == \"png\":\n",
    "        images.append(imageio.imread(filename))\n",
    "\n",
    "imageio.mimsave(\"Generating_Images.gif\", images)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Pro4.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}