generative_adversarial_networks.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "generative-adversarial-networks.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jarrydmartinx/generative-models/blob/master/generative_adversarial_networks.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ReTsAsRsiTN",
        "colab_type": "text"
      },
      "source": [
        "# Generative Adversarial Networks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YVGKGFHlsi9y",
        "colab_type": "text"
      },
      "source": [
        "## Install/Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fuhc4lAwpk7N",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 275
        },
        "outputId": "22dd8285-c125-4430-f6ad-75edf8d8da15"
      },
      "source": [
        "# @title Imports\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import sonnet as snt\n",
        "import pandas as pd\n",
        "import plotnine as gg\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<p style=\"color: red;\">\n",
              "The default version of TensorFlow in Colab will soon switch to TensorFlow 2.x.<br>\n",
              "We recommend you <a href=\"https://www.tensorflow.org/guide/migrate\" target=\"_blank\">upgrade</a> now \n",
              "or ensure your notebook will continue to use TensorFlow 1.x via the <code>%tensorflow_version 1.x</code> magic:\n",
              "<a href=\"https://colab.research.google.com/notebooks/tensorflow_version.ipynb\" target=\"_blank\">more info</a>.</p>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/util.py:63: The name tf.GraphKeys is deprecated. Please use tf.compat.v1.GraphKeys instead.\n",
            "\n",
            "WARNING:tensorflow:\n",
            "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
            "For more information, please see:\n",
            "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
            "  * https://github.com/tensorflow/addons\n",
            "  * https://github.com/tensorflow/io (for I/O related ops)\n",
            "If you depend on functionality not listed there, please file an issue.\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "je9TUwTSSdkq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# @title Get the MNIST data.\n",
        "mnist = tfds.load('mnist')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "prJFrsbx5Vfn",
        "colab_type": "text"
      },
      "source": [
        "## GAN training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R___KVXcQkLe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# @title Hyperparameters.\n",
        "\n",
        "BATCH_SIZE = 32\n",
        "MIXING_RATIO = 0.5  # Fraction of MNIST vs fake (generated) images.\n",
        "\n",
        "LATENT_SIZE = 12  # Size of the generator's latent space.\n",
        "GENERATOR_HIDDENS = [50, 50, 50]\n",
        "DISCRIMINATOR_HIDDENS = [50, 50]\n",
        "\n",
        "NUM_TRAINING_STEPS = 2000\n",
        "NUM_GENERATOR_STEPS_PER_DISCRIMINATOR_STEP = 10\n",
        "\n",
        "DISC_LEARNING_RATE = 1e-3\n",
        "GEN_LEARNING_RATE = 1e-3"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U5s9p2atqXS0",
        "colab_type": "code",
        "outputId": "a2da56b3-453c-4964-8626-f1921b76d2ef",
        "cellView": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "# @title Make & train a GAN.\n",
        "\n",
        "tf.reset_default_graph()\n",
        "\n",
        "mnist = tfds.load('mnist')\n",
        "\n",
        "dataset = mnist['train']\n",
        "\n",
        "# Given the dataset, create an iterator that produces MNIST images.\n",
        "real_images = (\n",
        "    dataset\n",
        "    .map(lambda x: tf.cast(x['image'], tf.float32) / tf.uint8.max)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .repeat()\n",
        "    .make_one_shot_iterator()\n",
        "    .get_next()\n",
        ")\n",
        "\n",
        "# Get image shape (e.g. for MNIST: (28, 28, 1)), set static batch size.\n",
        "image_shape = real_images.shape.as_list()[1:]  # [H, W, C]\n",
        "real_images.set_shape([BATCH_SIZE] + image_shape)  # [B, H, W, C]\n",
        "\n",
        "# Create the latent (noise) Tensor.\n",
        "latent = tf.random_normal(shape=(BATCH_SIZE, LATENT_SIZE))  # [B, L]\n",
        "\n",
        "# Create the generator network.\n",
        "generator = snt.Sequential([\n",
        "    snt.nets.MLP(GENERATOR_HIDDENS + [np.prod(image_shape)], \n",
        "                 name='generator'),\n",
        "    lambda x: tf.reshape(x, [BATCH_SIZE] + image_shape),\n",
        "#     tf.nn.sigmoid,\n",
        "])\n",
        "\n",
        "# This Tensor some fake images from the generator.\n",
        "fake_images = generator(latent)  # [B, H, W, C]\n",
        "\n",
        "# Mix real and fake images together randomly.\n",
        "real = tf.random_uniform(shape=(BATCH_SIZE,)) < MIXING_RATIO\n",
        "input_images = tf.where(real, real_images,  fake_images)\n",
        "\n",
        "# Create the discriminator network.\n",
        "discriminator = snt.Sequential([\n",
        "    snt.BatchFlatten(),\n",
        "    snt.nets.MLP(DISCRIMINATOR_HIDDENS + [2], name='discriminator'),\n",
        "])\n",
        "\n",
        "# The discriminator tries to classify inputs as either 'real' or 'fake'.\n",
        "logits = discriminator(input_images)\n",
        "\n",
        "# The discriminator loss is the binary cross-entropy.\n",
        "loss = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "    logits=logits, labels=tf.cast(real, tf.int32))\n",
        "loss = tf.reduce_mean(loss)\n",
        "\n",
        "# Create an op to track the discriminator's accuracy.\n",
        "predicted = tf.argmax(logits, axis=-1, output_type=tf.int32)\n",
        "acc = tf.cast(tf.equal(predicted, tf.cast(real, tf.int32)), tf.float32)\n",
        "acc = tf.reduce_mean(acc)\n",
        "\n",
        "# Create separate optimizers for the discriminator and generator.\n",
        "disc_optimizer = tf.train.AdamOptimizer(DISC_LEARNING_RATE)\n",
        "gen_optimizer = tf.train.AdamOptimizer(GEN_LEARNING_RATE)\n",
        "\n",
        "# Get the discriminator and generator network variables.\n",
        "disc_vars = discriminator.get_all_variables()\n",
        "gen_vars = generator.get_all_variables()\n",
        "\n",
        "# The discriminator seeks to *minimize* the discrimination loss.\n",
        "disc_grads_and_vars = disc_optimizer.compute_gradients(loss, disc_vars)\n",
        "disc_sgd_op = disc_optimizer.apply_gradients(disc_grads_and_vars)\n",
        "\n",
        "# The generator seeks to *maximize* the discrimination loss.\n",
        "gen_grads_and_vars = gen_optimizer.compute_gradients(-loss, gen_vars)\n",
        "gen_sgd_op = gen_optimizer.apply_gradients(gen_grads_and_vars)\n",
        "\n",
        "# Now we run our training loop.\n",
        "with tf.train.MonitoredSession() as sess:\n",
        "  results = []\n",
        "  for step in range(NUM_TRAINING_STEPS):\n",
        "    # Do a batch of SGD on the generator network.\n",
        "    sess.run(gen_sgd_op)\n",
        "    \n",
        "    # Periodically do a batch of SGD on the discriminator network.\n",
        "    if step % NUM_GENERATOR_STEPS_PER_DISCRIMINATOR_STEP == 0:\n",
        "      sess.run(disc_sgd_op)\n",
        "    \n",
        "    # Log the loss and discriminator accuracy.\n",
        "    loss_val, acc_val = sess.run([loss, acc])\n",
        "    results.append({'step': step,\n",
        "                    'loss': loss_val, \n",
        "                    'accuracy': acc_val})\n",
        "    if step % max(1, NUM_TRAINING_STEPS // 10) == 0:\n",
        "      print('Step {}/{}. Loss: {:.2f}. Acc: {:.2f}'.format(\n",
        "          step, NUM_TRAINING_STEPS, loss_val, acc_val))\n",
        "  \n",
        "  # When training is finished, generate a batch of samples.\n",
        "  samples = sess.run(fake_images)\n",
        "  \n",
        "  # Collect results into a Pandas dataframe for plotting.\n",
        "  df = pd.DataFrame(results)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:Entity <function <lambda> at 0x7fbe7a829f28> could not be transformed and will be executed as-is. Please report this to the AutoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: expected exactly one node node, found []\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:Entity <function <lambda> at 0x7fbe7a829f28> could not be transformed and will be executed as-is. Please report this to the AutoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: expected exactly one node node, found []\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING: Entity <function <lambda> at 0x7fbe7a829f28> could not be transformed and will be executed as-is. Please report this to the AutoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: expected exactly one node node, found []\n",
            "WARNING:tensorflow:From <ipython-input-4-42eab487b9c7>:12: DatasetV1.make_one_shot_iterator (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use `for ... in dataset:` to iterate over a dataset. If using `tf.estimator`, return the `Dataset` object directly from your input function. As a last resort, you can use `tf.compat.v1.data.make_one_shot_iterator(dataset)`.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From <ipython-input-4-42eab487b9c7>:12: DatasetV1.make_one_shot_iterator (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use `for ... in dataset:` to iterate over a dataset. If using `tf.estimator`, return the `Dataset` object directly from your input function. As a last resort, you can use `tf.compat.v1.data.make_one_shot_iterator(dataset)`.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:177: The name tf.make_template is deprecated. Please use tf.compat.v1.make_template instead.\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:177: The name tf.make_template is deprecated. Please use tf.compat.v1.make_template instead.\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:579: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:579: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:278: The name tf.get_default_graph is deprecated. Please use tf.compat.v1.get_default_graph instead.\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:278: The name tf.get_default_graph is deprecated. Please use tf.compat.v1.get_default_graph instead.\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/basic.py:126: calling TruncatedNormal.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/basic.py:126: calling TruncatedNormal.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/basic.py:131: calling Zeros.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/basic.py:131: calling Zeros.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/basic.py:247: The name tf.get_variable is deprecated. Please use tf.compat.v1.get_variable instead.\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/basic.py:247: The name tf.get_variable is deprecated. Please use tf.compat.v1.get_variable instead.\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From <ipython-input-4-42eab487b9c7>:38: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From <ipython-input-4-42eab487b9c7>:38: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:693: The name tf.get_collection is deprecated. Please use tf.compat.v1.get_collection instead.\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /tensorflow-1.15.0/python3.6/sonnet/python/modules/base.py:693: The name tf.get_collection is deprecated. Please use tf.compat.v1.get_collection instead.\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Graph was finalized.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Graph was finalized.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Running local_init_op.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Running local_init_op.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Done running local_init_op.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Done running local_init_op.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Step 0/2000. Loss: 0.65. Acc: 0.88\n",
            "Step 200/2000. Loss: 1153.76. Acc: 0.44\n",
            "Step 400/2000. Loss: 0.32. Acc: 0.62\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0atF-eJpz1FE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# @title Plot\n",
        "\n",
        "p = (gg.ggplot(df)\n",
        "     + gg.aes(x='step', y='accuracy')   \n",
        "     + gg.geom_line(color='Navy')\n",
        "     + gg.ggtitle('Discriminator Accuracy')\n",
        "    )\n",
        "p"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V01_75wvA0iR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# @title Show some samples (needs tuning/improvement!)\n",
        "plt.imshow(samples[4].squeeze())\n",
        "plt.grid(False)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}