AGP_AMallasto_WGAN_Tutorial.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/psuarezserrato/aprendizaje-geometrico/blob/main/AGP_AMallasto_WGAN_Tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5F181do92bq4"
      },
      "source": [
        "[Motivational video](https://www.youtube.com/watch?v=XOxxPcy5Gr4&t=81s)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lCljIveUruq5"
      },
      "source": [
        "#Brief Introduction to WGANs\n",
        "The aim of this tutorial is to give a brief introduction to Wasserstein GANs (WGANs). The objective is simple, we wish to learn from a target data distribution $\\mu_t$, by minimizing the $1$-Wasserstein distance $W_1$ between $\\mu_t$ and a parametrized model distribution $\\mu_\\omega$ with parameter $\\omega$. That is, our training objective is\n",
        "\\begin{equation}\n",
        "\\min\\limits_{\\omega} W_1(\\mu_{\\omega}, \\mu_t).\n",
        "\\end{equation}\n",
        "This has two main ingredients: 1) expressing the model distribution as a **push-forward** employing a **generator** neural network 2)  Estimating the $1$-Wasserstein distance using a **discriminator** neural network. Below, we will first briefly look at pytorch notation. After this, we consider 1) first, which is the easy part. Then, we move on to 2), where the magic of WGANs happens."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pGnnXyjOZlLA"
      },
      "source": [
        "#Import the python libraries that we will need\n",
        "\n",
        "import numpy as np #Standard python package for numerical computations\n",
        "\n",
        "import torch #GPU enabled deep learning library by Facebook\n",
        "import torch.nn as nn\n",
        "import torch.autograd as autograd\n",
        "from torch.autograd import Variable\n",
        "\n",
        "import seaborn as sns #For statistical visualization\n",
        "import pandas as pd #Cannot do data science without pandas\n",
        "\n",
        "import matplotlib.pyplot as plt #For visualization"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4FOT8edh44n3"
      },
      "source": [
        "# Some Pytorch Notation\n",
        "Pytorch operates on its own data structure, namely [torch.Tensors](https://pytorch.org/docs/stable/tensors.html). Below, we will define one, apply a function on it, and show  how pytorch can automatically compute gradients using [Autograd](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wXCsi1rA438x"
      },
      "source": [
        "#Define a tensor\n",
        "x = torch.Tensor([[1,2,3,4]])\n",
        "x = Variable(x, requires_grad = True)\n",
        "\n",
        "#Tell pytorch that we want to compute gradients\n",
        "x.requires_grad = True\n",
        "\n",
        "#Define a simple function using python's lambda notation\n",
        "f = lambda x: torch.sum(x**2)\n",
        "\n",
        "#Compute f(x)\n",
        "y = f(x)\n",
        "\n",
        "#Compute the gradient at x by backpropagation\n",
        "y.backward()\n",
        "v = x.grad\n",
        "\n",
        "print(v)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VkG6SXiA7Xsw"
      },
      "source": [
        "##Exercise 1:\n",
        "Now you do it! Compute the gradient of \n",
        "\\begin{equation}\n",
        "A \\mapsto \\log\\left(\\mathrm{det}(AB + A^{-1})\\right)\n",
        "\\end{equation}\n",
        "at \n",
        "\\begin{equation}\n",
        "A = \\begin{bmatrix}5&2\\\\2&1\\end{bmatrix},\n",
        "\\end{equation}\n",
        "when\n",
        "\\begin{equation}\n",
        "B = \\begin{bmatrix}1&2\\\\2&1\\end{bmatrix}.\n",
        "\\end{equation}\n",
        "( Hint: torch.log, torch.det, torch.inverse, matrix multiplication is given by A@B and for example the identity matrix is written as torch.Tensor([[1,0],[0,1]]) )"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c6HsfNeL7YYv"
      },
      "source": [],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h99lx4w9bjZ_"
      },
      "source": [
        "\n",
        "#  Generative Adversial Networks\n",
        "**Generative adversial networks ** (GANs) aim at learning a generative model for sampling from a given **data distribution** $\\mu_{\\mathrm{data}}\\in \\mathcal{P}(\\mathbb{R}^n)$. This is carried out by defining a **source distribution** $\\mu_{\\mathrm{source}}\\in \\mathcal{P}(\\mathbb{R}^d)$, where $d<<n$, and then pushing it forward with the **generator** $g_\\omega: \\mathbb{R}^d \\to \\mathbb{R}^n$, denoted by $(g_\\omega)_\\# \\mu_{\\mathrm{source}}$. Then, the parameter $\\omega$ is optimized to minimize a given similarity measure between  $(g_\\omega)_\\# \\mu_{\\mathrm{source}}$ and $\\mu_{\\mathrm{data}}$\n",
        "\n",
        "Two reminders before we move on. First, recall that the **manifold assumption** is prevalent in machine learning, stating that natural data lies on low dimensional submanifolds of the ambient data space, and therefore $d << n$ is justified. Second, the **push-forward** of a measure $\\mu$ on $\\Omega$ with respect to a measurable map $f:\\Omega \\to \\Omega'$ is defined by $(f_\\#\\mu)(A) = \\mu(f^{-1}(A))$ for any measurable set $A$ in $\\Omega'$. This definition  translates into something even simpler. Assume a random variable $X$ has law $\\mu$, then $f(X)$ has law $f_\\#\\mu$.\n",
        "\n",
        "# Implementing the Push-Forward\n",
        "Lets dirty our hands a bit, by creating a source distribution $\\mu_{\\mathrm{source}}\\in \\mathbb{R}^2$ and a generator $g_\\omega:\\mathbb{R}^2\\to \\mathbb{R}^2$ given by a multilayer perceptron (the simplest neural network), and see what we get.\n",
        "\n",
        "The function **source** eats the argument N and spits out N samples from the standard normal distribution on $\\mathbb{R^2}$.\n",
        "\n",
        "The generator **g** is constructed as a multilayer perceptron (MLP) with two hidden layers of 128 neurons, using rectified linear units (ReLU) as activation functions. This feeds on $M\\times 2$ matrices, applying itself row-wise, returning a  $M\\times 2$ matrix.\n",
        "\n",
        "See the [wikipedia article](https://en.wikipedia.org/wiki/Multilayer_perceptron) for more about MLPs.\n",
        "\n",
        "See the [pytorch documentation](https://pytorch.org/docs/stable/nn.html#non-linear-activations-weighted-sum-nonlinearity) for different activations and layers pytorch has to offer. Feel free to play around with them, if you are feeling adventurous!\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "32b4bYWEg06G"
      },
      "source": [
        "#Source samples [N] points from the standard normal distribution in R^2.\n",
        "#These will be given in a N x 2 matrix, with rows corresponding to samples.\n",
        "source = lambda N:torch.randn((N,2))\n",
        "#Generator network (used to push forward the source to the data distribution).\n",
        "g = nn.Sequential(nn.Linear(2,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,2)\n",
        "                   )\n",
        "\n",
        "#Sample 100 points from source and the push-forward.\n",
        "#We pick different samples for the push-forward to avoid correlation.\n",
        "samples_source = source(100)\n",
        "samples_push = g(source(100)).detach()\n",
        "#above, detach tells pytorch that this variable is now static, and should not\n",
        "#be considered in the computational graph. This allows us to plot the points.\n",
        "\n",
        "#Lets see how the push-forward changes the distribution.\n",
        "plt.figure()\n",
        "plt.scatter(samples_source[:,0], samples_source[:,1], color='b')\n",
        "plt.scatter(samples_push[:,0], samples_push[:,1], color='r')\n",
        "plt.legend(['source', 'push-forward'])\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lMgcqX2ep6Is"
      },
      "source": [
        "# Picking a Similarity Measure\n",
        "Right now, our push-forward is not producing anything useful. To learn the model to sample from a given data distribution $\\mu_{\\mathrm{data}}$, we need to pick a similarity measure to be minimized. This choice specifies which GAN we are working with, for example,  [The original GAN](https://papers.nips.cc/paper/5423-generative-adversarial-nets) minimizes the [Jensen-Shannon divergence](https://en.wikipedia.org/wiki/Jensen%E2%80%93Shannon_divergence), whereas the [WGAN](https://arxiv.org/abs/1701.07875) minimizes the [$1$-Wasserstein metric](https://en.wikipedia.org/wiki/Wasserstein_metric). There are plenty of other choices, but in this tutorial we will focus on WGANs.\n",
        "\n",
        "#$p$-Wasserstein Metric\n",
        "Let us define the $p$-Wasserstein metric between two probability measures, that results from solving a constrained linear program. We first give the **primal** formulation, after which we look at the **dual formulation**.\n",
        "\n",
        "##Primal formulation\n",
        "Let $(X,d)$ be a metric space that is complete and separable (a Polish space), and pick two probability measures $\\mu$ and $\\nu$ with finite $p$-moments. Then, the $p$-Wasserstein metric between $\\mu$ and $\\nu$ is given by\n",
        "\\begin{equation}\n",
        "W_p(\\mu,\\nu) = \\left(\\min_\\gamma\\mathbb{E}_\\gamma[d^p]\\right)^{\\frac{1}{p}},\n",
        "\\end{equation}\n",
        "where $\\gamma$ is constrained to be a joint distribution of $\\mu$ and $\\nu$, and we use the notation\n",
        "\\begin{equation}\n",
        "\\mathbb{E}_\\mu[f] = \\int_X f(x)d\\mu(x).\n",
        "\\end{equation}\n",
        "Note that this defines a symmetric and positive-definite metric that satisfies the triangle inequality.\n",
        "\n",
        "##Dual formulation\n",
        "As the Wasserstein metric results from solving a linear program, it admits a dual formulation given by \n",
        "\\begin{equation}\n",
        "W_p^p(\\mu, \\nu) = \\max_{\\varphi, \\psi}\\lbrace \\mathbb{E}_\\mu[\\varphi] + \\mathbb{E}_\\nu[\\psi]\\rbrace,\n",
        "\\end{equation}\n",
        "where we require $\\varphi(x) + \\psi(y) \\leq d^p(x,y)$ for any $x,y$. The optimal $\\varphi, \\psi$ are called **Kantorovich potentials**, satisfying $\\varphi(x) + \\psi(y) = d^p(x,y)$ for any $(x,y)$ in the support of the optimal $\\gamma$ that solves the primal problem. Furthermore, we know that $\\psi(y) = \\varphi^c(y)$, where\n",
        "\\begin{equation}\n",
        "\\varphi^c(y) = \\inf_x\\{d^p(x,y)-\\varphi(x)\\},\n",
        "\\end{equation}\n",
        "is the **$c$-transform** of $\\varphi$. This allows us to write the dual problem as\n",
        "\\begin{equation}\n",
        "W_p^p(\\mu, \\nu) = \\max_{\\varphi, \\psi}\\lbrace \\mathbb{E}_\\mu[\\varphi] + \\mathbb{E}_\\nu[\\varphi^c]\\rbrace.\n",
        "\\end{equation}\n",
        "If $\\varphi$ is $1$-Lipschitz and $c=d^1$, magic happens, as $\\varphi^c(y) = -\\varphi(y)$!(not factorial) It can be shown, that when $p=1$, the optimal $\\varphi$ will indeed be $1$-Lipschitz, which we will use later on when formulating the WGAN objective. (Trivia: for a general $p$, if the support of any joint measure of $\\mu$ and $\\nu$ is bounded by diameter $D$, then the Kantorovich potentials will be $pD^{p-1}$-Lipschitz. However, the $d^p$-transform does not behave as nicely for Lipschitz functions).\n",
        "\n",
        "#Computing the 1-Wasserstein Distance\n",
        "Let's put this into practice and compute the 1-Wasserstein distance between $\\mu_{\\mathrm{source}}$ and $(g_\\omega)_\\#\\mu_{\\mathrm{source}}$, which we implemented earlier. We will do this using a Monte Carlo scheme and approximating $\\varphi$ with a MLP $f_{\\omega'}$, which we call the **discriminator**.\n",
        "\n",
        "Sample $N$ points from the source $\\mu_s$ and push-forward $(g_\\omega)_\\#\\mu_s$, yielding $\\{x_i\\}_{i=1}^N$ and $\\{y_i\\}_{i=1}^N$, respectively, which we call the **mini-batches**. Then, we compute the objective function $\\rho$\n",
        "\\begin{equation}\n",
        "\\rho(\\omega') = W_1(\\mu_{\\mathrm{source}}, (g_\\omega)_\\#\\mu_{\\mathrm{source}}) \\approx \\max_{\\omega'} \\frac{1}{N}\\lbrace \\sum_{i=1}^N f_{\\omega'}(x_i) - \\sum_{i=1}^N f{\\omega'}(y_i)  \\rbrace,\n",
        "\\end{equation}\n",
        "where we used that at optimality, $f_{\\omega'}^c(y) = -f_{\\omega'}(y)$. We optimize using [stochastic gradient descent](https://en.wikipedia.org/wiki/Stochastic_gradient_descent), which means that after sampling mini-batches, we compute the gradient $v$ of the above expression with respect to $\\omega'$, update $v + \\omega' \\mapsto \\omega'$, and then sample new mini-batches and repeat the same until convergence (or until we are satisfied).\n",
        "\n",
        "Computing the gradients and updating the parameters has been made very easy in pytorch, by using automatic differentation via [backpropagation](https://en.wikipedia.org/wiki/Backpropagation), which is carried out by [pytorch optimizers](https://pytorch.org/docs/stable/optim.html).\n",
        "\n",
        "##Implementation\n",
        "The implementation can be found below. We define the discriminator and the related optimizer (1), after which we begin the training loop (2). At each step in the training loop, we sample mini-batches from each measure (3), and compute the mean of $f_{\\omega'}$ under $\\mu$ and the mean of $-f_{\\omega'}$ under $\\nu$ (4). Summing these together gives us the objective function (5). Then, we just compute the gradient of the loss with respect to $\\omega'$ and update $\\omega'$ (6)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m7O0IDW11-TW"
      },
      "source": [
        "\n",
        "#-----------(1)-----------\n",
        "#Define the discriminator that assings a real value to\n",
        "#points in R^2.\n",
        "f = nn.Sequential(nn.Linear(2,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,1)\n",
        "                   )\n",
        "#For updating the parameters, we need to use a pytorch\n",
        "#optimizer object.\n",
        "f_optim = torch.optim.RMSprop(f.parameters(), lr = 1e-4)\n",
        "#We are using the RMSprop optimizer, but see the above cell\n",
        "#for other optimization methods pytorch has to offer.\n",
        "\n",
        "N_batch = 64 #Mini-batch size\n",
        "N_iterations = 1000 #Amount of gradient steps we will take\n",
        "\n",
        "#Lets keep track of how our discriminator is approximating\n",
        "#the 1-Wasserstein distance at each iteration\n",
        "loss_history = []\n",
        "\n",
        "#-----------(2)-----------\n",
        "#Train the discriminator\n",
        "for i in range(N_iterations):\n",
        "  #-----------(3)-----------\n",
        "  #Sample from the source and the push-forward\n",
        "  samples_source = source(N_batch)\n",
        "  samples_push = g(source(N_batch))\n",
        "  \n",
        "  #-----------(4)-----------\n",
        "  #Compute f values\n",
        "  f_source = f(samples_source)\n",
        "  f_push = -f(samples_push)\n",
        "  \n",
        "  #-----------(5)-----------\n",
        "  #Pytorch optimizers minimize objective functions, as we\n",
        "  #want to maximize instead, we will be minimizing the negative\n",
        "  #of the objective\n",
        "  loss = -(f_source.mean() + f_push.mean())\n",
        "  \n",
        "  #-----------(6)-----------\n",
        "  #tell pytorch to compute the gradient\n",
        "  loss.backward()\n",
        "  \n",
        "  #update the parameters of f\n",
        "  f_optim.step()\n",
        "  \n",
        "  #Pytorch accummulates gradients, so we will have to zero\n",
        "  #them after each update\n",
        "  f_optim.zero_grad()\n",
        "  \n",
        "  #save the current value\n",
        "  loss_history.append(float(-loss))\n",
        "  \n",
        "#Plot the loss\n",
        "plt.figure()\n",
        "plt.plot(np.arange(N_iterations), loss_history)\n",
        "plt.title('Approximating 1-Wasserstein')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "reekse1qDEW6"
      },
      "source": [
        "##Exercise 2: \n",
        "Increase N_iterations and observe the behavior of the loss function. Are we able to compute the 1-Wasserstein distance?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JdvI0ifI-yfD"
      },
      "source": [
        "#Incorporating the Discriminator Constraints\n",
        "As you can see, the loss does not converge. This is because we are not taking into account the constraint on the Kantorovich potentials\n",
        "\\begin{equation}\n",
        "\\varphi(x)-\\varphi(y) \\leq d(x,y).\n",
        "\\end{equation}\n",
        "The constraint is very important, indeed, as the only way the $l^2$-metric $d$ comes into play is through the constraint. Therefore, in the cell above, we are not computing the 1-Wasserstein distance!\n",
        "\n",
        "Implementing the constraint is an art of its own, as a considerable amount of papers study how to do this properly. The original WGAN paper does this through **weight clipping**. Other notable versions include the [WGAN-GP](https://arxiv.org/abs/1704.00028), which is pretty much the state of the art, and the [CT- WGAN](https://arxiv.org/abs/1803.01541).\n",
        "\n",
        "The papers approach the constraint through Lipschitzness. If you look at the expression above, this is exactly the requirement for $\\varphi$ to be $1$-Lipschitz. The papers then focus on enforcing this:\n",
        "\n",
        "-Original WGANs clip the weights of the neural network $f_{\\omega'}$ to a box (they enforce $-c\\leq \\omega'_{\\mathrm{weights}} \\leq c$, where $c$ is a small constant and the inequalities are considered element-wise). This forces $f_{\\omega'}$ to be $K$-Lipschitz for some $K$, which is not a problem, as maximizing the dual expression with respect to $K$-Lipschitz functions yields the Wasserstein distance multiplied by $K$.\n",
        "\n",
        "-Gradient penalty WGAN (WGAN-GP) relies on the result, that $f_{\\omega'}$ is $1$-Lipschitz, if $\\|\\nabla_x f_{\\omega'}(x)\\| \\leq 1$ for almost every $x$. Therefore, they apply a penalty term to the objective function to enforce this condition.\n",
        "\n",
        "-Consistency term WGAN (CT-WGAN) enforces the Lipschitz condition directly, by adding a penalty term $\\mathbb{E}(|f_{\\omega'}(x) - f_{\\omega'}(y)|-1)^2$.\n",
        "\n",
        "Do note, that these techniques only work in the $1$-Wasserstein case. Below, we will apply the original weight clipping scheme with $c=0.01$. This does not perform too well in practice, but is the easiest to implement. Afterwards, we will try out the gradient penalty used in WGAN-GP."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Zh6KFIHGuBB"
      },
      "source": [
        "##Exercise 3:\n",
        "I am lying when saying that we are not taking the constraint on Kantorovich potentials into accout. Why? \n",
        "\n",
        "**Extra**: We are taking the constraint into account in a way. Why isn't the method converging?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B-lpUoqkNG19"
      },
      "source": [
        "#An utility function for clipping the weights of f\n",
        "def clip_weights(f, c):\n",
        "  for param in f.parameters():\n",
        "    if len(param)>1:\n",
        "        param.data.clamp_(-c, c)\n",
        "\n",
        "#Discriminator\n",
        "f = nn.Sequential(nn.Linear(2,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,1)\n",
        "                   )\n",
        "f_optim = torch.optim.RMSprop(f.parameters(), lr = 1e-4)\n",
        "\n",
        "N_batch = 64 #Mini-batch size\n",
        "N_iterations = 1000 #Amount of gradient steps we will take\n",
        "\n",
        "loss_history = [] #Save the loss at each iteration here\n",
        "\n",
        "#Train the discriminator\n",
        "for i in range(N_iterations):\n",
        "    #Sample from the source and the push-forward\n",
        "    samples_source = source(N_batch)\n",
        "    samples_push = g(source(N_batch))\n",
        "\n",
        "    #-----------Weight clipping-----------\n",
        "    #Clip the weights of f\n",
        "    clip_weights(f, .01) \n",
        "    \n",
        "    #Compute the f values\n",
        "    f_source = f(samples_source)\n",
        "    f_push = -f(samples_push)\n",
        "    loss = -(f_source.mean() + f_push.mean())\n",
        "\n",
        "    #Compute the gradient\n",
        "    loss.backward()\n",
        "\n",
        "    #update the parameters of f\n",
        "    f_optim.step()\n",
        "    f_optim.zero_grad()\n",
        "\n",
        "    #save the current value\n",
        "    loss_history.append(float(-loss))\n",
        "  \n",
        "#Plot the loss\n",
        "plt.figure()\n",
        "plt.plot(np.arange(N_iterations), loss_history)\n",
        "plt.title('Approximating 1-Wasserstein with Weight Clipping')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "371-wUJwShbj"
      },
      "source": [
        "# Gradient Penalty\n",
        "We are doing better, but the objective still will not converge in a reasonable amount of iterations. Let's see, how the gradient penalty introduced in WGAN-GP performs. This is done, by adding the gradient penalty to the objective function\n",
        "\\begin{equation}\n",
        "\\rho(\\omega') = W_1(\\mu_{\\mathrm{source}}, (g_{\\omega})_\\#\\mu_{\\mathrm{source}}) \\approx \\max_{\\omega'} \\frac{1}{N}\\lbrace \\sum_{i=1}^N f_{\\omega'}(x_i) - \\sum_{i=1}^N f_{\\omega'}(y_i) - \\lambda\\mathbb{E}_\\nu[(\\|\\nabla_x \\varphi (x)\\| -1)^2]\\rbrace,\n",
        "\\end{equation}\n",
        "where $\\lambda$ is the weight of the penalization, and $\\nu$ some reference measure. Note that in theory, instead of equality we should be enforcing $\\|\\nabla_x f_{\\omega'}(x)\\| \\leq 1$, but in the WGAN-GP the authors remarked that in practice this suffices."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L-5gr8ZLHNPe"
      },
      "source": [
        "##Exercise 4:\n",
        "\n",
        "Below, I have provided a function that computes the gradient penalty. Incorporate this into the discriminator training procedure."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AkLCsxukSivt"
      },
      "source": [
        "#An utility function for clipping the weights of f\n",
        "def gradient_penalty(f, samples_source, samples_push, lambda_reg=5, use_cuda = False):\n",
        "    BATCH_SIZE = samples_source.shape[0]\n",
        "    LAMBDA = 5\n",
        "    alpha = torch.rand(BATCH_SIZE, 1)\n",
        "    alpha = alpha.expand(samples_source.size())\n",
        "\n",
        "    interpolates = alpha * samples_source + ((1 - alpha) * samples_push)\n",
        "    \n",
        "    if use_cuda:\n",
        "        interpolates = interpolates.cuda(gpu)\n",
        "    interpolates = autograd.Variable(interpolates, requires_grad=True)\n",
        "\n",
        "    disc_interpolates = f(interpolates)\n",
        "\n",
        "    gradients = autograd.grad(outputs=disc_interpolates, inputs=interpolates,\n",
        "                              grad_outputs=torch.ones(disc_interpolates.size()),\n",
        "                              create_graph=True, retain_graph=True, only_inputs=True)[0]\n",
        "\n",
        "    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * lambda_reg\n",
        "    return gradient_penalty\n",
        "\n",
        "#Discriminator\n",
        "f = nn.Sequential(nn.Linear(2,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,1)\n",
        "                   )\n",
        "f_optim = torch.optim.RMSprop(f.parameters(), lr = 1e-4)\n",
        "\n",
        "N_batch = 64 #Mini-batch size\n",
        "N_iterations = 1000 #Amount of gradient steps we will take\n",
        "\n",
        "loss_history = [] #Save the loss at each iteration here\n",
        "\n",
        "#Train the discriminator\n",
        "for i in range(N_iterations):\n",
        "    #Sample from the source and from the push-forward\n",
        "    samples_source = source(N_batch)\n",
        "    samples_push = g(source(N_batch))\n",
        "    \n",
        "   #Remember to add the gradient penalty!!\n",
        "        \n",
        "    f_source = f(samples_source)\n",
        "    f_push = -f(samples_push)\n",
        "    loss = -(f_source.mean() + f_push.mean())\n",
        "     \n",
        "    #backpropagate for the gradient\n",
        "    loss.backward()\n",
        "\n",
        "    #update the parameters of f\n",
        "    f_optim.step()\n",
        "    f_optim.zero_grad()\n",
        "\n",
        "    #save the current value\n",
        "    loss_history.append(float(-loss))\n",
        "  \n",
        "#Plot the loss\n",
        "plt.figure()\n",
        "plt.plot(np.arange(N_iterations), loss_history)\n",
        "plt.title('Approximating 1-Wasserstein with Gradient Penalty')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YF2e1ej6Z5wr"
      },
      "source": [
        "#Implementing the WGAN\n",
        "We can now estimate the $1$-Wasserstein distance between two measures $\\mu$ and $\\nu$, where above we took $\\mu=\\mu_s$ and $\\nu = (g_{\\omega})_\\#\\mu_s$.  The final step in the WGAN implementation is to update the weights $\\omega$ of the generator to minimize the objective function that is given by the $1$-Wasserstein distance and possibly a penalization term.\n",
        "\n",
        "To do this, we will do something slightly more interesting, and try to learn the distribution of a Gaussian mixture model with three components in $\\mathbb{R}^2$, which we have implemented below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R1fHGh4GaBYr"
      },
      "source": [
        "def gaussian_mixture(N, weights, variances, means):\n",
        "    M = len(weights) #Number of clusters\n",
        "    d = means[0].shape[1] #Dimension of the Gaussian\n",
        "    \n",
        "    #Sample from multinomial distribution\n",
        "    pvals = torch.multinomial(weights, N, replacement = True)\n",
        "    \n",
        "    #counts of each index appearing in pvals\n",
        "    counts = [int(torch.sum(pvals == i)) for i in range(M)]\n",
        "    \n",
        "    #Sample from clusters\n",
        "    samples = []\n",
        "    \n",
        "    for i in range(M):\n",
        "        samplesi = torch.sqrt(variances[i])*torch.randn((counts[i],d)) + means[i]\n",
        "        samples.append(samplesi)\n",
        "        \n",
        "    return torch.cat(samples,0)\n",
        "    \n",
        "weights = torch.Tensor([.2, .3, .5])\n",
        "variances = torch.Tensor([.5, .2, .6])\n",
        "means = [torch.Tensor([[0,2]]), torch.Tensor([[-4,-1]]), torch.Tensor([[2,-2]])]\n",
        "\n",
        "#Sample data points for visualization\n",
        "samples = gaussian_mixture(200, weights, variances, means)\n",
        "\n",
        "\n",
        "#visualize data\n",
        "plt.figure()\n",
        "plt.scatter(samples[:,0], samples[:,1])\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FuC8vlSvfDBF"
      },
      "source": [
        "#Updating the Generator\n",
        "Now that we have our data set, we will learn how to sample from it."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Q_afj3_b8Xd"
      },
      "source": [
        "\n",
        "#Discriminator\n",
        "f = nn.Sequential(nn.Linear(2,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,1)\n",
        "                   )\n",
        "f_optim = torch.optim.RMSprop(f.parameters(), lr = 1e-3)\n",
        "\n",
        "#Source\n",
        "source = lambda N:torch.randn((N,2))\n",
        "#Target\n",
        "target = lambda N: gaussian_mixture(N, weights, variances, means)\n",
        "\n",
        "#Generator\n",
        "g = nn.Sequential(nn.Linear(2,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,2)\n",
        "                   )\n",
        "g_optim = torch.optim.RMSprop(g.parameters(), lr = 1e-3)\n",
        "\n",
        "N_batch = 64 #Mini-batch size\n",
        "N_discriminator = 5 #Discriminator iterations per generator iteration\n",
        "N_generator = 500 #Amount of generator iterations\n",
        "\n",
        "loss_history = []\n",
        "#Train the generator\n",
        "for i in range(N_generator):\n",
        "    #Sample from the source and apply the push-forward\n",
        "    samples_source= source(N_batch)\n",
        "    #Sample from the target\n",
        "    #samples_nu = target(N_batch)\n",
        "    samples_nu = target(N_batch)\n",
        "    \n",
        "    #Train the discriminator\n",
        "    for j in range(N_discriminator):\n",
        "        #Push-forward source samples\n",
        "        samples_mu = g(samples_source)\n",
        "        \n",
        "        #Compute the gradient penalty\n",
        "        penalty = gradient_penalty(f, samples_mu, samples_nu)\n",
        "        \n",
        "        #Compute f values\n",
        "        f_mu = f(samples_mu)\n",
        "        f_nu = -f(samples_nu)\n",
        "        loss_discriminator = -(f_mu.mean() + f_nu.mean()) + penalty\n",
        "\n",
        "        #backpropagate for the gradient\n",
        "        loss_discriminator.backward()\n",
        "\n",
        "        #update the parameters of f\n",
        "        f_optim.step()\n",
        "        f_optim.zero_grad()\n",
        "        g_optim.zero_grad()\n",
        "\n",
        "        #save the current value\n",
        "        loss_history.append(float(-loss_discriminator))\n",
        "   \n",
        "    #Update generator\n",
        "    samples_mu = g(samples_source)\n",
        "    f_source = f(samples_mu)\n",
        "    f_push = -f(samples_nu)\n",
        "    \n",
        "    #This time we want to minimize the expression, so no minus needed.\n",
        "    loss_generator = f_source.mean() + f_push.mean()\n",
        "    \n",
        "    loss_generator.backward()\n",
        "    g_optim.step()\n",
        "    g_optim.zero_grad()\n",
        "    \n",
        "#Lets plot the results!\n",
        "samples_generator = g(source(200)).detach()\n",
        "samples_target = target(200).detach()\n",
        "\n",
        "gen_samples = g(source(200)).detach()\n",
        "dat_samples =pd.DataFrame(target(200).detach().numpy(), columns=['x', 'y'])\n",
        "gen_samples_df = pd.DataFrame(gen_samples.detach().numpy(), columns=['x', 'y'])\n",
        "\n",
        "plt.figure()\n",
        "plt.scatter(dat_samples.x, dat_samples.y, c='b', alpha=0.2)\n",
        "sns.kdeplot(gen_samples_df.x, gen_samples_df.y, zorder=0, n_levels=10, legend=True, shade=True)\n",
        "plt.legend(['Target Samples'])\n",
        "plt.title('Model distribution')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79Kow8EFOCFe"
      },
      "source": [
        "#A Challenge for the Brave\n",
        "Below, I have added a script for downloading the MNIST hand-written digits dataset. Running the script will yield a **train_loader**, that is mandatory for working with large datasets. It reads mini-batches from the data into memory when needed, and discards then when we are done with that mini-batch. This avoids loading the whole dataset into memory at once. Below this script, there is a cell with the skeleton for the WGAN procedure on the MNIST dataset. It is your task to complete this.\n",
        "\n",
        "In order to to make this a bit more interesting, we will use the GPU provided by google for this task. To do this, we will have to do some notifications in the code. We will have to change our datatype from torch.FloatTensor to torch.cuda.FloatTensor. Additionally, we will also have to specify that our discriminator and generator work with this datatype, by using the command .cuda() (see the code below).\n",
        "\n",
        "Changing the notebook to use the GPU will reset the session, so we will have to import the libraries again below. You can start using the GPU, by clicking on [Runtime] in the upper left corner, then choose [Change runtime type] and finally pick GPU as the [Hardware accelerator]."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MqEwmUigbEtB"
      },
      "source": [
        "#Import the python libraries that we will need\n",
        "\n",
        "import numpy as np #Standard python package for numerical computations\n",
        "\n",
        "import torch #GPU enabled deep learning library by Facebook\n",
        "import torch.nn as nn\n",
        "import torch.autograd as autograd\n",
        "from torch.autograd import Variable\n",
        "import torchvision.datasets as dset\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "import seaborn as sns #For statistical visualization\n",
        "import pandas as pd #Cannot do data science without pandas\n",
        "\n",
        "import matplotlib.pyplot as plt #For visualization\n",
        "\n",
        "def imshow(image):\n",
        "    I = image.clone().cpu().numpy().reshape(DIM_IMG, DIM_IMG) + 0.5\n",
        "    plt.imshow(I, cmap='gray')\n",
        "\n",
        "DLATENT = 64\n",
        "DIM_IMG = 28\n",
        "DIM = DIM_IMG*DIM_IMG\n",
        "N_batch = 64\n",
        "trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (.5,))])\n",
        "# if not exist, download mnist dataset\n",
        "train_set = dset.MNIST(root = './', train=True, transform=trans, download=True)\n",
        "\n",
        "\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(\n",
        "                 dataset=train_set,\n",
        "                 batch_size=N_batch,\n",
        "                 shuffle=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfLidQJJ06kI"
      },
      "source": [
        "##WGAN Skeleton for the MNIST Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ow2bZ42zORSj"
      },
      "source": [
        "#We have to work with a GPU friendly data type.\n",
        "dtype = torch.cuda.FloatTensor\n",
        "\n",
        "def gradient_penalty(f, samples_source, samples_push, lambda_reg=5, use_cuda = False):\n",
        "    BATCH_SIZE = samples_source.shape[0]\n",
        "    LAMBDA = 5\n",
        "    alpha = torch.rand(BATCH_SIZE, 1).type(dtype)\n",
        "    alpha = alpha.expand(samples_source.size())\n",
        "\n",
        "    interpolates = alpha * samples_source + ((1 - alpha) * samples_push)\n",
        "    \n",
        "    if use_cuda:\n",
        "        interpolates = interpolates.cuda(gpu)\n",
        "    interpolates = autograd.Variable(interpolates, requires_grad=True)\n",
        "\n",
        "    disc_interpolates = f(interpolates)\n",
        "\n",
        "    gradients = autograd.grad(outputs=disc_interpolates, inputs=interpolates,\n",
        "                              grad_outputs=torch.ones(disc_interpolates.size()).type(dtype),\n",
        "                              create_graph=True, retain_graph=True, only_inputs=True)[0]\n",
        "\n",
        "    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * lambda_reg\n",
        "    return gradient_penalty\n",
        "  \n",
        "  \n",
        "#Discriminator\n",
        "f = nn.Sequential(nn.Linear(DIM,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,1)\n",
        "                   ).cuda() #<--- Notice the .cuda()\n",
        "\n",
        "f_optim = torch.optim.Adam(f.parameters(), lr = 1e-4, betas=(0, .9))\n",
        "\n",
        "#Source\n",
        "source = lambda N:torch.randn((N,DLATENT)).type(dtype)\n",
        "\n",
        "#Generator\n",
        "g = nn.Sequential(nn.Linear(DLATENT,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,128), nn.ReLU(),\n",
        "                    nn.Linear(128,DIM), nn.Tanh()\n",
        "                   ).cuda()\n",
        "g_optim = torch.optim.Adam(g.parameters(), lr = 1e-4, betas=(0,.9))\n",
        "\n",
        "N_batch = 64 #Mini-batch size\n",
        "N_discriminator = 5 #Discriminator iterations per generator iteration\n",
        "N_epoch= 10 #Amount of times we iterate over the full dataset\n",
        "\n",
        "loss_history = []\n",
        "#Train the generator\n",
        "for i in range(N_epoch):\n",
        "  #Sample from target using the train_loader\n",
        "    for samples_nu, _ in train_loader:\n",
        "        if samples_nu.shape[0] != N_batch:\n",
        "          break\n",
        "\n",
        "        samples_nu = samples_nu.view(N_batch,-1).type(dtype) #Make the data GPU friendly\n",
        "        samples_source= source(N_batch)\n",
        "\n",
        "        #Train the discriminator\n",
        "        for j in range(N_discriminator):\n",
        "            #Add the code!\n",
        "\n",
        "        #Update generator\n",
        "        #Add the code!\n",
        "\n",
        "    print('%d Epochs done' %i)\n",
        "    \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hFPzvjdcyRTd"
      },
      "source": [
        "#Plot some generator samples!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C2TXhR3ER1ic"
      },
      "source": [
        "samples_generator = g(source(8)).detach()\n",
        "\n",
        "fig = plt.figure(figsize=(8,4))\n",
        "for i in range(8):\n",
        "  plt.subplot(2,4, i+1)\n",
        "  imshow(samples_generator[i])\n",
        "  plt.grid(False)\n",
        "  plt.axis('off')\n",
        "  plt.xticks([])\n",
        "  plt.yticks([])\n",
        "  plt.subplots_adjust(wspace=0, hspace=0, left=0, right=1, bottom=0, top=1)\n",
        "plt.show()\n",
        "print(torch.max(samples_generator[0]))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}