GANs_with_TF_2_0.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GANs with TF 2.0.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/TF-2.0-Hacks/blob/master/GANs_with_TF_2_0.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s9vaOlDrphRN",
        "colab_type": "text"
      },
      "source": [
        "This notebook follows [this amazing tutorial on GANs](https://medium.com/@devnag/generative-adversarial-networks-gans-in-50-lines-of-code-pytorch-e81b79659e3f) and tries to port the code to TensorFlow 2.0. \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qn636Xcba-Q1",
        "colab_type": "text"
      },
      "source": [
        "## Install `Tensorflow 2.0`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qp5esV6aANHl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install tensorflow-gpu==2.0.0-beta1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6jIXCCVsbC-K",
        "colab_type": "text"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jGB-ehhCAsO0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from tensorflow import keras\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7EUQq5NbFSe",
        "colab_type": "text"
      },
      "source": [
        "## Helper function to generate a distribution (normal) for the real data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7G6PMW-RBTpC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_distribution_sampler(mu, sigma):\n",
        "  return lambda n: tf.convert_to_tensor(np.random.normal(mu, sigma, (1, n)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "loyhvfRxbLX2",
        "colab_type": "text"
      },
      "source": [
        "## Helper function to generate a uniform distribution for the generator network"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NzGT3Q8zBzht",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_generator_input_sampler():\n",
        "  return lambda m, n: tf.convert_to_tensor(np.random.rand(m, n))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZG6Py0Xmbczy",
        "colab_type": "text"
      },
      "source": [
        "## The Generator network class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4ifARasTCMw0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Generator(keras.Model):\n",
        "  def __init__(self, input_size, hidden_size, output_size):\n",
        "    super(Generator, self).__init__()\n",
        "    self.map1 = keras.layers.Dense(hidden_size, input_shape=input_size, activation='tanh')\n",
        "    self.map2 = keras.layers.Dense(hidden_size, activation='tanh')\n",
        "    self.map3 = keras.layers.Dense(output_size, activation='linear')\n",
        "  \n",
        "  def call(self, inputs):\n",
        "    x = self.map1(inputs)\n",
        "    x = self.map2(x)\n",
        "    x = self.map3(x)\n",
        "    return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DZSWuojAbgDE",
        "colab_type": "text"
      },
      "source": [
        "## The Discriminator network class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FAAzJHCFEp9e",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Discriminator(keras.Model):\n",
        "  def __init__(self, input_size, hidden_size, output_size):\n",
        "    super(Discriminator, self).__init__()\n",
        "    self.map1 = keras.layers.Dense(hidden_size, input_shape=input_size, activation='sigmoid')\n",
        "    self.map2 = keras.layers.Dense(hidden_size, activation='sigmoid')\n",
        "    self.map3 = keras.layers.Dense(output_size, activation='sigmoid')\n",
        "  \n",
        "  def call(self, inputs):\n",
        "    x = self.map1(inputs)\n",
        "    return self.map3(self.map2(x))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9-PY9BZuFgqr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_moments(d):\n",
        "  # https://stats.stackexchange.com/questions/126346/why-kurtosis-of-a-normal-distribution-is-3-instead-of-0\n",
        "  # Return the first 4 moments of the data provided\n",
        "  d = tf.transpose(d, (1, 0))\n",
        "  mean = tf.reduce_mean(d)\n",
        "  diffs = (d - mean)\n",
        "  var = tf.reduce_mean(tf.pow(diffs, 2.0))\n",
        "  std = tf.sqrt(var)\n",
        "  zscores = diffs / std\n",
        "  skews = tf.reduce_mean(tf.pow(zscores, 3.0))\n",
        "  kurtoses = tf.reduce_mean(tf.pow(zscores, 4.0)) - 3.0 # excess kurtosis, should be 0 for Gaussian\n",
        "  return tf.stack([mean, std, skews, kurtoses], axis=0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3OV4KWgRJy7i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def stats(d):\n",
        "    return [np.mean(d), np.std(d)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WzKl011RXlLU",
        "colab_type": "text"
      },
      "source": [
        "## Model hyperparameters and other constants"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-qj45CReOW9c",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Model parameters\n",
        "g_input_size = 1      # Random noise dimension coming into generator, per output vector\n",
        "g_hidden_size = 5     # Generator complexity\n",
        "g_output_size = 1     # Size of generated output vector\n",
        "d_input_size = 500    # Minibatch size - cardinality of distributions\n",
        "d_hidden_size = 10    # Discriminator complexity\n",
        "d_output_size = 1     # Single dimension for 'real' vs. 'fake' classification\n",
        "minibatch_size = d_input_size\n",
        "\n",
        "d_learning_rate = 1e-3\n",
        "g_learning_rate = 1e-3\n",
        "sgd_momentum = 0.9\n",
        "\n",
        "num_epochs = 5000\n",
        "print_interval = 100\n",
        "d_steps = 20\n",
        "g_steps = 20\n",
        "\n",
        "dfe, dre, ge = 0, 0, 0\n",
        "d_real_data, d_fake_data, g_fake_data = None, None, None"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OdRJkiDAXwQk",
        "colab_type": "text"
      },
      "source": [
        "## Data generation parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CM9ZeUMqOt8j",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data_mean = 4\n",
        "data_stddev = 1.25\n",
        "\n",
        "d_sampler = get_distribution_sampler(data_mean, data_stddev)\n",
        "gi_sampler = get_generator_input_sampler()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UIYpY99oXyiQ",
        "colab_type": "text"
      },
      "source": [
        "## Initialize the networks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8LN_PKnbO4Ag",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "G = Generator(input_size=(500, 1),\n",
        "                  hidden_size=g_hidden_size,\n",
        "                  output_size=g_output_size)\n",
        "\n",
        "D = Discriminator(input_size=(1,4),\n",
        "                  hidden_size=d_hidden_size,\n",
        "                  output_size=d_output_size)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FWvzGpJcX2AQ",
        "colab_type": "text"
      },
      "source": [
        "## Declare the loss and optimizers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lkc9zVjCPlaW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "criterion = tf.keras.losses.BinaryCrossentropy(from_logits=True)  \n",
        "d_optimizer = tf.keras.optimizers.SGD(learning_rate=d_learning_rate, momentum=sgd_momentum)\n",
        "g_optimizer = tf.keras.optimizers.SGD(learning_rate=g_learning_rate, momentum=sgd_momentum)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TeBmj8g9X8FV",
        "colab_type": "text"
      },
      "source": [
        "## One forward and backward pass with the Discriminator network with real data\n",
        "\n",
        "*We do not update the parameters with these gradients.*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Kw8MU7Vgr1E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# d_real_data = d_sampler(d_input_size)\n",
        "\n",
        "# with tf.GradientTape() as tape:\n",
        "#   d_real_decision = D(get_moments(d_real_data).reshape((1,4)))\n",
        "#   d_real_error = criterion(d_real_decision, np.ones((1,1)))  # ones = true\n",
        "# d_real_grads = tape.gradient(d_real_error, D.trainable_weights) # compute/store gradients, but don't change params\n",
        "# d_real_grads[0].numpy()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "81YBUAh6YGc2",
        "colab_type": "text"
      },
      "source": [
        "## One forward and backward pass with the Discriminator network with the fake data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IjcWC0gNibld",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# d_gen_input = gi_sampler(minibatch_size, g_input_size)\n",
        "# with tf.GradientTape() as tape:\n",
        "#   with tape.stop_recording():\n",
        "#     d_fake_data = G(d_gen_input)\n",
        "#   d_fake_decision = D(get_moments(d_fake_data.numpy().T).reshape((1,4)))\n",
        "#   d_fake_error = criterion(d_fake_decision, np.zeros((1,1)))\n",
        "# d_fake_grads = tape.gradient(d_fake_error, D.trainable_weights) \n",
        "# print(d_fake_grads[0].numpy())\n",
        "# d_optimizer.apply_gradients(zip(d_fake_grads, D.trainable_weights)) # Only optimizes D's parameters"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dS7mg238Yf3M",
        "colab_type": "text"
      },
      "source": [
        "## One forward and backward pass with the Generator network"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yp3eWwowQsfA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# gen_input = gi_sampler(minibatch_size, g_input_size)\n",
        "# with tf.GradientTape() as tape:\n",
        "#   g_fake_data = G(gen_input)\n",
        "#   dg_fake_decision = D(tf.reshape(get_moments_tf(g_fake_data), (1, 4)))\n",
        "#   g_error = criterion(dg_fake_decision, np.ones((1,1)))\n",
        "# g_grads = tape.gradient(g_error, G.trainable_weights)\n",
        "# g_optimizer.apply_gradients(zip(g_grads, G.trainable_weights))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "as0XOQuhQVai",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 496
        },
        "outputId": "fc0c169d-9b00-407d-e97d-de07469d4147"
      },
      "source": [
        "for epoch in range(num_epochs):\n",
        "    for d_index in range(d_steps):\n",
        "        #  1A: Train D on real\n",
        "        d_real_data = tf.convert_to_tensor(d_sampler(d_input_size))\n",
        "        with tf.GradientTape() as tape:\n",
        "          d_real_decision = D(tf.reshape(get_moments(d_real_data), (1,4)))\n",
        "          d_real_error = criterion(d_real_decision, tf.convert_to_tensor(np.ones((1,1))))  # ones = true\n",
        "        d_real_grads = tape.gradient(d_real_error, D.trainable_weights) # compute/store gradients, but don't change params\n",
        "        \n",
        "        #  1B: Train D on fake\n",
        "        d_gen_input = tf.convert_to_tensor(gi_sampler(minibatch_size, g_input_size))\n",
        "        with tf.GradientTape() as tape:\n",
        "          with tape.stop_recording():\n",
        "            d_fake_data = G(d_gen_input)\n",
        "          d_fake_decision = D(tf.reshape(get_moments(d_fake_data), (1, 4)))\n",
        "          d_fake_error = criterion(d_fake_decision, tf.convert_to_tensor(np.zeros((1,1))))\n",
        "        d_fake_grads = tape.gradient(d_fake_error, D.trainable_weights) \n",
        "        d_optimizer.apply_gradients(zip(d_fake_grads, D.trainable_weights)) # Only optimizes D's parameters\n",
        "\n",
        "        dre, dfe = d_real_error.numpy(), d_fake_error.numpy()\n",
        "\n",
        "    for g_index in range(g_steps):\n",
        "        # 2. Train G on D's response (but DO NOT train D on these labels)\n",
        "        gen_input = tf.convert_to_tensor(gi_sampler(minibatch_size, g_input_size))\n",
        "        with tf.GradientTape() as tape:\n",
        "          g_fake_data = G(gen_input)\n",
        "          dg_fake_decision = D(tf.reshape(get_moments(g_fake_data), (1, 4)))\n",
        "          g_error = criterion(dg_fake_decision, tf.convert_to_tensor(np.ones((1,1))))\n",
        "        g_grads = tape.gradient(g_error, G.trainable_weights)\n",
        "        g_optimizer.apply_gradients(zip(g_grads, G.trainable_weights))\n",
        "\n",
        "        ge = g_error.numpy()\n",
        "\n",
        "    if epoch % print_interval == 0:\n",
        "        print(\"Epoch %s: D (%s real_err, %s fake_err) G (%s err); Real Dist (%s),  Fake Dist (%s) \" %\n",
        "              (epoch, dre, dfe, ge, stats(d_real_data.numpy()), stats(d_fake_data.numpy())))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING: Logging before flag parsing goes to stderr.\n",
            "W0825 16:46:38.575693 140618269017984 deprecation.py:323] From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/nn_impl.py:182: add_dispatch_support.<locals>.wrapper (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": [
            "Epoch 0: D (1.0886667966842651 real_err, 0.6931471824645996 fake_err) G (1.05702543258667 err); Real Dist ([3.969248032721655, 1.2093934306513305]),  Fake Dist ([-0.16539627014683225, 0.07419386361400125]) \n",
            "Epoch 100: D (1.087760329246521 real_err, 0.6931471824645996 fake_err) G (1.0470203161239624 err); Real Dist ([4.066130571911628, 1.2346445557152335]),  Fake Dist ([-0.30376876966542266, 0.11773658791402217]) \n",
            "Epoch 200: D (1.0887233018875122 real_err, 0.6931471824645996 fake_err) G (1.0439269542694092 err); Real Dist ([3.932937238259309, 1.2277380207905655]),  Fake Dist ([-0.45581279916255424, 0.05699663380595091]) \n",
            "Epoch 300: D (1.089498519897461 real_err, 0.6931471824645996 fake_err) G (1.0407992601394653 err); Real Dist ([3.9936716043970866, 1.26235733971329]),  Fake Dist ([-0.7447241589696177, 0.008942774214919355]) \n",
            "Epoch 400: D (1.0871673822402954 real_err, 0.6931471824645996 fake_err) G (1.0361649990081787 err); Real Dist ([3.984654450017385, 1.2337837894636985]),  Fake Dist ([-1.0351879749603823, 0.002109290445755961]) \n",
            "Epoch 500: D (1.0883467197418213 real_err, 0.6931471824645996 fake_err) G (1.0363883972167969 err); Real Dist ([3.988741519040915, 1.2769630954792424]),  Fake Dist ([-1.3187153942607712, 0.01051046466078226]) \n",
            "Epoch 600: D (1.0887354612350464 real_err, 0.6931471824645996 fake_err) G (1.0328069925308228 err); Real Dist ([4.0354432685602495, 1.2639294650160453]),  Fake Dist ([-1.5836209810851478, 0.0016828689300378383]) \n",
            "Epoch 700: D (1.086708903312683 real_err, 0.6931471824645996 fake_err) G (1.0311508178710938 err); Real Dist ([4.031299864056426, 1.2657557941081317]),  Fake Dist ([-1.8416030141383084, 0.0013018275670126536]) \n",
            "Epoch 800: D (1.0865145921707153 real_err, 0.6931471824645996 fake_err) G (1.0310169458389282 err); Real Dist ([3.932973583748997, 1.2239113147645544]),  Fake Dist ([-2.0886074957547605, 0.0013422930905573522]) \n",
            "Epoch 900: D (1.0907894372940063 real_err, 0.6931471824645996 fake_err) G (1.0300602912902832 err); Real Dist ([3.9778720529350466, 1.2423837728686176]),  Fake Dist ([-2.3256808724775326, 0.0012080414429308768]) \n",
            "Epoch 1000: D (1.0880186557769775 real_err, 0.6931471824645996 fake_err) G (1.028603434562683 err); Real Dist ([3.9940469689561717, 1.2715260343372974]),  Fake Dist ([-2.551267113414703, 0.0011714463221389963]) \n",
            "Epoch 1100: D (1.0890711545944214 real_err, 0.6931471824645996 fake_err) G (1.0281318426132202 err); Real Dist ([3.9846537807237072, 1.2359541446670275]),  Fake Dist ([-2.7675533651296775, 0.0012208285143596883]) \n",
            "Epoch 1200: D (1.08739173412323 real_err, 0.6931471824645996 fake_err) G (1.0269633531570435 err); Real Dist ([3.9974220963831475, 1.1644718718455738]),  Fake Dist ([-2.9766773087855567, 0.0011700479420714103]) \n",
            "Epoch 1300: D (1.0886634588241577 real_err, 0.6931471824645996 fake_err) G (1.0266574621200562 err); Real Dist ([4.083303671966571, 1.2616810053521563]),  Fake Dist ([-3.1772018743392696, 0.0014295534361247573]) \n",
            "Epoch 1400: D (1.08675217628479 real_err, 0.6931471824645996 fake_err) G (1.0262341499328613 err); Real Dist ([3.9615456319438453, 1.2818510111661996]),  Fake Dist ([-3.369814808413284, 0.0013407713700283883]) \n",
            "Epoch 1500: D (1.087263584136963 real_err, 0.6931471824645996 fake_err) G (1.025339961051941 err); Real Dist ([3.904134046604551, 1.1628845651645383]),  Fake Dist ([-3.5559914304462503, 0.001369007176916181]) \n",
            "Epoch 1600: D (1.0901052951812744 real_err, 0.6931471824645996 fake_err) G (1.025075912475586 err); Real Dist ([4.059169023803737, 1.2646959132035673]),  Fake Dist ([-3.7347559002639072, 0.0014854473813596168]) \n",
            "Epoch 1700: D (1.0903997421264648 real_err, 0.6931471824645996 fake_err) G (1.0246912240982056 err); Real Dist ([4.073631580767418, 1.307118740402538]),  Fake Dist ([-3.9077235078747, 0.00145042718194546]) \n",
            "Epoch 1800: D (1.0890417098999023 real_err, 0.6931471824645996 fake_err) G (1.0242844820022583 err); Real Dist ([4.013299261717157, 1.2009866954408044]),  Fake Dist ([-4.074154845765458, 0.0014768626935643397]) \n",
            "Epoch 1900: D (1.0875056982040405 real_err, 0.6931471824645996 fake_err) G (1.0239284038543701 err); Real Dist ([3.932749537333192, 1.3076516555178603]),  Fake Dist ([-4.234803080252499, 0.0016178279077750752]) \n",
            "Epoch 2000: D (1.0874297618865967 real_err, 0.6931471824645996 fake_err) G (1.0234853029251099 err); Real Dist ([3.9801396846719146, 1.2030596669445992]),  Fake Dist ([-4.3905574960095075, 0.0014821050846723775]) \n",
            "Epoch 2100: D (1.0854185819625854 real_err, 0.6931471824645996 fake_err) G (1.023307204246521 err); Real Dist ([3.9527717265787685, 1.2955137559221834]),  Fake Dist ([-4.5413157765176875, 0.0013579439905201547]) \n",
            "Epoch 2200: D (1.0870475769042969 real_err, 0.6931471824645996 fake_err) G (1.022774577140808 err); Real Dist ([4.032350958520261, 1.2474040559753414]),  Fake Dist ([-4.6874410387605545, 0.00129060017867163]) \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H9ggO7OcwgMV",
        "colab_type": "text"
      },
      "source": [
        "The network is still not properly configured. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zB1GPjYvT8SY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}