DCGAN_Mnist_Image_Generator.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyPMMIYlUnZFCYTD08+YW80Q",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/masoudshahrian/DeepLearning-Code/blob/main/DCGAN_Mnist_Image_Generator.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ITcNPzlG2zbi"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras.layers import Dense, Reshape, Flatten, Dropout, LeakyReLU, BatchNormalization, Conv2D, Conv2DTranspose\n",
        "from tensorflow.keras.datasets import mnist\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Loading the MNIST dataset\n",
        "(x_train, _), (_, _) = mnist.load_data()\n",
        "x_train = (x_train - 127.5) / 127.5 # Normalize the images to [-1, 1]\n",
        "x_train = np.expand_dims(x_train, axis=-1)\n",
        "\n",
        "BUFFER_SIZE = 60000\n",
        "BATCH_SIZE = 256"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UtlJvbh020ln",
        "outputId": "6329222c-e487-4bf0-861a-d9da7b01a183"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
            "\u001b[1m11490434/11490434\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 0us/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Create batches of the training data\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices(x_train).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)\n",
        "\n",
        "# Generator model\n",
        "def make_generator_model():\n",
        "    model = tf.keras.Sequential()\n",
        "    model.add(Dense(7*7*256, use_bias=False, input_shape=(100,)))\n",
        "    model.add(BatchNormalization())\n",
        "    model.add(LeakyReLU())\n",
        "    model.add(Reshape((7, 7, 256)))\n",
        "    model.add(Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))\n",
        "    model.add(BatchNormalization())\n",
        "    model.add(LeakyReLU())\n",
        "    model.add(Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))\n",
        "    model.add(BatchNormalization())\n",
        "    model.add(LeakyReLU())\n",
        "    model.add(Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))\n",
        "    return model\n",
        "\n",
        "# Discriminator model\n",
        "def make_discriminator_model():\n",
        "    model = tf.keras.Sequential()\n",
        "    model.add(Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[28, 28, 1]))\n",
        "    model.add(LeakyReLU())\n",
        "    model.add(Dropout(0.3))\n",
        "    model.add(Conv2D(128, (5, 5), strides=(2, 2), padding='same'))\n",
        "    model.add(LeakyReLU())\n",
        "    model.add(Dropout(0.3))\n",
        "    model.add(Flatten())\n",
        "    model.add(Dense(1))\n",
        "    return model\n",
        "\n",
        "# Define the loss and optimizers\n",
        "cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
        "\n",
        "def discriminator_loss(real_output, fake_output):\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",
        "    return total_loss\n",
        "\n",
        "def generator_loss(fake_output):\n",
        "    return cross_entropy(tf.ones_like(fake_output), fake_output)\n",
        "\n",
        "generator = make_generator_model()\n",
        "discriminator = make_discriminator_model()\n",
        "\n",
        "generator_optimizer = tf.keras.optimizers.Adam(1e-4)\n",
        "discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)"
      ],
      "metadata": {
        "id": "j1oNM69k2_TO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "# Define the training loop\n",
        "EPOCHS = 50\n",
        "noise_dim = 100\n",
        "num_examples_to_generate = 16\n",
        "\n",
        "# Seed for visualization\n",
        "seed = tf.random.normal([num_examples_to_generate, noise_dim])\n",
        "\n",
        "@tf.function\n",
        "def train_step(images):\n",
        "    noise = tf.random.normal([BATCH_SIZE, noise_dim])\n",
        "    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
        "        generated_images = generator(noise, training=True)\n",
        "        real_output = discriminator(images, training=True)\n",
        "        fake_output = discriminator(generated_images, training=True)\n",
        "        gen_loss = generator_loss(fake_output)\n",
        "        disc_loss = discriminator_loss(real_output, fake_output)\n",
        "    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)\n",
        "    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n",
        "    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))\n",
        "    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))\n",
        "\n",
        "def train(dataset, epochs):\n",
        "    for epoch in range(epochs):\n",
        "        for image_batch in dataset:\n",
        "            train_step(image_batch)\n",
        "        display.clear_output(wait=True)\n",
        "        generate_and_save_images(generator, epoch + 1, seed)\n",
        "        print(f'Epoch {epoch+1} completed')\n",
        "\n",
        "def generate_and_save_images(model, epoch, test_input):\n",
        "    predictions = model(test_input, training=False)\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, :, :, 0] * 127.5 + 127.5, cmap='gray')\n",
        "        plt.axis('off')\n",
        "    plt.savefig(f'image_at_epoch_{epoch}.png')\n",
        "    plt.show()\n",
        "\n",
        "# Start training\n",
        "train(train_dataset, EPOCHS)\n"
      ],
      "metadata": {
        "id": "IO89qLJL21OG"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}