Normalizing_flows_empty_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Normalizing_flows_empty.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "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.3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mlelarge/dataflowr/blob/master/Normalizing_flows_empty_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NA1jh8NDbTSz"
      },
      "source": [
        "# Normalizing flows"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7r7vD1Ft2hc6",
        "colab_type": "text"
      },
      "source": [
        "Here is very good read on normalizing flows : [blogpost](https://lilianweng.github.io/lil-log/2018/10/13/flow-based-deep-generative-models.html)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "onzQnZoFbTS0"
      },
      "source": [
        "# Density estimation using Real NVP\n",
        "\n",
        "Normalizing flows is the class of probabilistic models that provides flexible parametrical probabilistic models, where the probability density function can be computed exactly. In the assignment, we will consider a real-valued non-volume preserving normalizing flows (Real NVP) -- a special case of normalizing flow.\n",
        "\n",
        "## Problem setting\n",
        "\n",
        "Our goal is to train a generative network $g_\\theta:  Z \\rightarrow X, g = f^{-1}$ that maps latent variable $z \\sim p(z)$ to a sample $x \\sim p(x)$. Where $p(z)$ is a prior distibiution and $p(x)$ is a data distibution. An illustrative example is provided below.\n",
        "\n",
        "![alt text](https://senya-ashukha.github.io/assignments/normalizing-flows/2d-example.png)\n",
        "\n",
        "## Change of variable formula\n",
        "\n",
        "Given an observed data variable $x \\in X$,\n",
        "a simple prior probability distribution $p_{Z}$ on a latent variable $z \\in Z$,\n",
        "and a bijection $f: X \\rightarrow Z$ (with $g = f^{-1}$),\n",
        "the change of variable formula defines a model distribution of $X$ by\n",
        "<!-- \\begin{align}\n",
        "p_{X}(x) &= p_{Z}\\big(f(x)\\big) \\left|\\det\\left(\\cfrac{\\partial f(x)}{\\partial x^T} \\right)\\right|\n",
        "\\label{eq:change-variables}\\\\\n",
        "\\log\\left(p_{X}(x)\\right) &= \\log\\Big(p_{Z}\\big(f(x)\\big)\\Big) + \\log\\left(\\left|\\det\\left(\\frac{\\partial f(x)}{\\partial x^T}\\right)\\right|\\right)\n",
        ",\n",
        "\\end{align} -->\n",
        "\n",
        "![alt text](https://senya-ashukha.github.io/assignments/normalizing-flows/f.png)\n",
        "where $\\frac{\\partial f(x)}{\\partial x^T}$ is the Jacobian of $f$ at $x$.\n",
        "\n",
        "Exact samples from the resulting distribution can be generated by using the inverse transform sampling rule. A sample $z \\sim p_{Z}$ is drawn in the latent space, and its inverse image $x = f^{-1}(z) = g(z)$ generates a sample in the original space. Computing the density at a point $x$ is accomplished by computing the density of its image $f(x)$ and multiplying by the associated Jacobian determinant $\\det\\left(\\frac{\\partial f(x)}{\\partial x^T}\\right)$.\n",
        "\n",
        "## Real NVP\n",
        "\n",
        "Real NVP presents a class of functions where $\\log\\left(\\left|\\det\\left(\\frac{\\partial f(x)}{\\partial x^T}\\right)\\right|\\right)$ can be computed efficiently (see, 3.3 Properties, https://arxiv.org/abs/1605.08803). Every layer of Real NVP is a coupling layer followed by permutation layer. Combination of coupling and permutation layers can be implemented as a masked version of the coupling layer:\n",
        "## $$y = b \\odot x + (1 - b) \\odot \\Big(x \\odot \\exp\\big(s(b \\odot x)\\big) + t(b \\odot x)\\Big)$$\n",
        "\n",
        "where $s$ and $t$ stand for scale and translation, and are functions from $R^{D} \\mapsto R^{D}$, and $\\odot$ is the Hadamard product or element-wise product, $b$ is a binary mask. For more details on the model see the paper Density estimation using Real NVP https://arxiv.org/abs/1605.08803."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NgMBEYdwbTS2"
      },
      "source": [
        "# Implementation of Real NVP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "loxXeNy9bTS3",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "\n",
        "from pylab import rcParams\n",
        "rcParams['figure.figsize'] = 5, 4\n",
        "rcParams['figure.dpi'] = 150\n",
        "\n",
        "import torch\n",
        "from torch import nn\n",
        "from torch import distributions\n",
        "from torch.nn.parameter import Parameter\n",
        "\n",
        "from sklearn import datasets"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iFaqlfn3HcOV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "nets = lambda: nn.Sequential(nn.Linear(2, 256), nn.LeakyReLU(), nn.Linear(256, 256), nn.LeakyReLU(), nn.Linear(256, 2), nn.Tanh())\n",
        "nett = lambda: nn.Sequential(nn.Linear(2, 256), nn.LeakyReLU(), nn.Linear(256, 256), nn.LeakyReLU(), nn.Linear(256, 2))\n",
        "# functions that take no arguments and return a pytorch model, dim(X) -> dim(X)\n",
        "\n",
        "masks = torch.from_numpy(np.array([[0, 1], [1, 0], [0, 1], [1, 0], [0, 1], [1, 0]]).astype(np.float32))\n",
        "# torch.Tensor of size #number_of_coupling_layers x #dim(X)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nFVr-ttwHcOX",
        "colab_type": "code",
        "outputId": "cc585ee6-2740-487b-96fd-919507c45b5f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 87
        }
      },
      "source": [
        "from torch import distributions\n",
        "prior = distributions.MultivariateNormal(torch.zeros(2), torch.eye(2))\n",
        "print(prior.log_prob(torch.Tensor([0,0])))\n",
        "print(prior.sample((3,)))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor(-1.8379)\n",
            "tensor([[ 0.8131, -0.4593],\n",
            "        [-1.0915,  2.1433],\n",
            "        [-0.3912, -0.0338]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "GpyZ5AqObTS7",
        "colab": {}
      },
      "source": [
        "class RealNVP(nn.Module):\n",
        "    def __init__(self, nets, nett, mask, prior):\n",
        "        super(RealNVP, self).__init__()\n",
        "        \n",
        "        # Create a flow\n",
        "        # nets:  a function that returns a PyTorch neural network, e.g., nn.Sequential, s = nets(), s: dim(X) -> dim(X)\n",
        "        # nett:  a function that returns a PyTorch neural network, e.g., nn.Sequential, t = nett(), t: dim(X) -> dim(X)\n",
        "        # mask:  a torch.Tensor of size #number_of_coupling_layers x #dim(X)\n",
        "        # prior: an object from torch.distributions e.g., torch.distributions.MultivariateNormal\n",
        "        \n",
        "        self.prior = prior\n",
        "        self.mask = nn.Parameter(mask, requires_grad=False)\n",
        "        self.t = torch.nn.ModuleList([nett() for _ in range(len(masks))])\n",
        "        self.s = torch.nn.ModuleList([nets() for _ in range(len(masks))])\n",
        "        \n",
        "    def g(self, z):\n",
        "        # Compute and return g(z) = x, \n",
        "        #    where self.mask[i], self.t[i], self.s[i] define a i-th masked coupling layer   \n",
        "        # z: a torch.Tensor of shape batchSize x 1 x dim(X)\n",
        "        # return x: a torch.Tensor of shape batchSize x 1 x dim(X)\n",
        "        return x\n",
        "\n",
        "    def f(self, x):        \n",
        "        # Compute f(x) = z and log_det_Jakobian of f, \n",
        "        #    where self.mask[i], self.t[i], self.s[i] define a i-th masked coupling layer   \n",
        "        # x: a torch.Tensor, of shape batchSize x dim(X), is a datapoint\n",
        "        # return z: a torch.Tensor of shape batchSize x dim(X), a hidden representations\n",
        "        # return log_det_J: a torch.Tensor of len batchSize\n",
        "        \n",
        "        return z, log_det_J\n",
        "    \n",
        "    def log_prob(self, x):\n",
        "        # Compute and return log p(x)\n",
        "        # using the change of variable formula and log_det_J computed by f\n",
        "        # return logp: torch.Tensor of len batchSize\n",
        "        return logp\n",
        "        \n",
        "    def sample(self, batchSize): \n",
        "        # Draw and return batchSize samples from flow using implementation of g\n",
        "        # return x: torch.Tensor of shape batchSize x 1 x dim(X)\n",
        "        return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "J6YPmM_3bTTE",
        "colab": {}
      },
      "source": [
        "flow = RealNVP(nets, nett, masks, prior)\n",
        "# Check that a flow is reversible g(f(x)) = x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "vb80JOSSbTTG",
        "colab": {}
      },
      "source": [
        "from sklearn import datasets\n",
        "\n",
        "trainable_parametrs = # list of all trainable parameters in a flow\n",
        "optimizer = # choose an optimizer, use module torch.optim\n",
        "\n",
        "for t in range(5001):    \n",
        "    noisy_moons = datasets.make_moons(n_samples=100, noise=.05)[0].astype(np.float32)\n",
        "    loss = # compute the maximum-likelihood loss\n",
        "    \n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    \n",
        "    if t % 500 == 0:\n",
        "        print('iter %s:' % t, 'loss = %.3f' % loss)\n",
        "        \n",
        "# Check that the loss decreases\n",
        "# Is the visualization below good?"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6hOxCpXYbTTJ"
      },
      "source": [
        "# Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "RmK9g7CIbTTK",
        "colab": {}
      },
      "source": [
        "from sklearn import datasets\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "from pylab import rcParams\n",
        "rcParams['figure.figsize'] = 10, 8\n",
        "rcParams['figure.dpi'] = 300"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Gz5YG8FNbTTN"
      },
      "source": [
        "Draw several plots: \n",
        "- samples from flow\n",
        "- samples from prior\n",
        "- data samples\n",
        "- mapping form data to prior\n",
        "\n",
        "The goal is to obtain figure similar to https://arxiv.org/abs/1605.08803"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "qbhwTtHjpHC9",
        "colab": {}
      },
      "source": [
        "noisy_moons = datasets.make_moons(n_samples=1000, noise=.05)[0].astype(np.float32)\n",
        "z = flow.f(torch.from_numpy(noisy_moons))[0].detach().numpy()\n",
        "plt.subplot(221)\n",
        "plt.scatter(z[:, 0], z[:, 1])\n",
        "plt.title(r'$z = f(X)$')\n",
        "\n",
        "z = np.random.multivariate_normal(np.zeros(2), np.eye(2), 1000)\n",
        "plt.subplot(222)\n",
        "plt.scatter(z[:, 0], z[:, 1])\n",
        "plt.title(r'$z \\sim p(z)$')\n",
        "\n",
        "plt.subplot(223)\n",
        "x = datasets.make_moons(n_samples=1000, noise=.05)[0].astype(np.float32)\n",
        "plt.scatter(x[:, 0], x[:, 1], c='r')\n",
        "plt.title(r'$X \\sim p(X)$')\n",
        "\n",
        "plt.subplot(224)\n",
        "x = flow.sample(1000).detach().numpy()\n",
        "plt.scatter(x[:, 0, 0], x[:, 0, 1], c='r')\n",
        "plt.title(r'$X = g(z)$')"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}