Module9/Normalizing_flows_empty.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NA1jh8NDbTSz"
   },
   "source": [
    "# Normalizing flows\n",
    "\n",
    "The image below is taken form this very good blog post on normalizing flows : [blogpost](https://lilianweng.github.io/lil-log/2018/10/13/flow-based-deep-generative-models.html)\n",
    "\n",
    "![](https://raw.githubusercontent.com/dataflowr/website/master/modules/extras/flows/three-generative-models.png)\n",
    "\n",
    "A **flow-based generative model** is constructed by a sequence of **invertible** transformations. The main advantage of flows is that the model explicitly learns the data distribution $p(\\mathbf{x})$ and therefore the loss function is simply the negative log-likelihood.\n",
    "\n",
    "Given a sample $\\mathbf{x}$ and a prior $p(\\mathbf{z})$, we compute $f(\\mathbf{x}) = \\mathbf{z}$ with an invertible function $f$ that will be learned. Given $f$ and the prior $p(\\mathbf{z})$, we can compute the evidence $p(\\mathbf{x})$ thanks to the change of variable formula:\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\mathbf{z} &\\sim p(\\mathbf{z}), \\mathbf{z} = f(\\mathbf{x}), \\\\\n",
    "p(\\mathbf{x}) \n",
    "&= p(\\mathbf{z}) \\left\\vert \\det \\dfrac{d \\mathbf{z}}{d \\mathbf{x}} \\right\\vert  \n",
    "= p(f(\\mathbf{x})) \\left\\vert \\det \\dfrac{\\partial f(\\mathbf{x})}{\\partial \\mathbf{x}} \\right\\vert\n",
    "\\end{align*}\n",
    "$$\n",
    "where $\\dfrac{\\partial f(\\mathbf{x})}{\\partial \\mathbf{x}}$ is the Jacobian matrix of $f$.\n",
    "Recall that given a function mapping a $n$-dimensional input vector $\\mathbf{x}$ to a $m$-dimensional output vector, $f: \\mathbb{R}^n \\mapsto \\mathbb{R}^m$, the matrix of all first-order partial derivatives of this function is called the **Jacobian matrix**, $J_f$ where one entry on the i-th row and j-th column is $(J_f(\\mathbf{x}))_{ij} = \\frac{\\partial f_i(\\mathbf{x})}{\\partial x_j}$:\n",
    "$$\n",
    "{J_f(\\mathbf{x})} = \\begin{bmatrix}\n",
    "\\frac{\\partial f_1(\\mathbf{x})}{\\partial x_1} & \\dots & \\frac{\\partial f_1(\\mathbf{x})}{\\partial x_n} \\\\[6pt]\n",
    "\\vdots & \\ddots & \\vdots \\\\[6pt]\n",
    "\\frac{\\partial f_m(\\mathbf{x})}{\\partial x_1} & \\dots & \\frac{\\partial f_m(\\mathbf{x})}{\\partial x_n} \\\\[6pt]\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "Below, we will parametrize $f$ with a neural network and learn $f$ by maximizing $\\ln p(\\mathbf{x})$. More precisely, given a dataset $(\\mathbf{x}_1,\\dots,\\mathbf{x}_n)$ and a model provided by a prior $p(\\mathbf{z})$ and a neural network $f$, we optimize the weights of $f$ by minimizing:\n",
    "$$\n",
    "-\\sum_{i}\\ln p(\\mathbf{x_i}) = \\sum_i -\\ln p(f(\\mathbf{x}_i)) -\\ln\\left\\vert \\det \\dfrac{\\partial f(\\mathbf{x}_i)}{\\partial \\mathbf{x}} \\right\\vert.\n",
    "$$\n",
    "\n",
    "**We need to ensure that $f$ is always invertible and that the determinant is simple to compute.**\n",
    "\n",
    "## [Density estimation using Real NVP](https://arxiv.org/abs/1605.08803) \n",
    "by Laurent Dinh, Jascha Sohl-Dickstein, Samy Bengio (2016)\n",
    "\n",
    "[Real NVP](https://arxiv.org/abs/1605.08803) uses function $f$ obtained by stacking affine coupling layers which for an input $\\mathbf{x}\\in \\mathbb{R}^D$ produce the output $\\mathbf{y}\\in\\mathbb{R}^D$ defined by ($d<D$): \n",
    "$$\n",
    "\\begin{align}\n",
    "\\label{eq:aff}\n",
    "\\mathbf{y}_{1:d} &= \\mathbf{x}_{1:d}\\\\\n",
    "\\mathbf{y}_{d+1:D} &= \\mathbf{x}_{d+1:D} \\odot \\exp\\left(s(\\mathbf{x}_{1:d})\\right) +t(\\mathbf{x}_{1:d}) ,\n",
    "\\end{align}\n",
    "$$\n",
    "where $s$ (scale) and $t$ (translation) are neural networks mapping $\\mathbb{R}^d$ to $\\mathbb{R}^{D-d}$ and $\\odot$ is the element-wise product.\n",
    "\n",
    "For any functions $s$ and $t$, the affine coupling layer is invertible:\n",
    "\\begin{align*}\n",
    "\\begin{cases}\n",
    "\\mathbf{y}_{1:d} &= \\mathbf{x}_{1:d} \\\\ \n",
    "\\mathbf{y}_{d+1:D} &= \\mathbf{x}_{d+1:D} \\odot \\exp({s(\\mathbf{x}_{1:d})}) + t(\\mathbf{x}_{1:d})\n",
    "\\end{cases}\n",
    "\\Leftrightarrow \n",
    "\\begin{cases}\n",
    "\\mathbf{x}_{1:d} &= \\mathbf{y}_{1:d} \\\\ \n",
    "\\mathbf{x}_{d+1:D} &= (\\mathbf{y}_{d+1:D} - t(\\mathbf{y}_{1:d})) \\odot \\exp(-s(\\mathbf{y}_{1:d}))\n",
    "\\end{cases}\n",
    "\\end{align*}\n",
    "\n",
    "The Jacobian of an affine coupling layer is a lower triangular matrix:\n",
    "\\begin{align*}\n",
    "J(\\mathbf{x}) =  \\frac{\\partial \\mathbf{y}}{\\partial \\mathbf{x}}=\n",
    "\\begin{bmatrix}\n",
    "  \\mathbb{I}_d & \\mathbf{0}_{d\\times(D-d)} \\\\[5pt]\n",
    "  \\frac{\\partial \\mathbf{y}_{d+1:D}}{\\partial \\mathbf{x}_{1:d}} & \\text{diag}(\\exp(s(\\mathbf{x}_{1:d})))\n",
    "\\end{bmatrix}\n",
    "\\end{align*}\n",
    "Hence the determinant is simply the product of terms on the diagonal:\n",
    "\\begin{align*}\n",
    "\\left\\vert\\det(J(\\mathbf{x}))\\right\\vert\n",
    "= \\prod_{j=1}^{D-d}\\exp(s(\\mathbf{x}_{1:d}))_j\n",
    "= \\exp\\left(\\sum_{j=1}^{D-d} s(\\mathbf{x}_{1:d})_j\\right)\n",
    "\\end{align*}\n",
    "Note that, we do not need to compute the Jacobian of $s$ or $t$ and to compute $f^{-1}$, we do not need to compute the inverse of $s$ or $t$ (which might not exist!). In other words, we can take arbitrary complex functions for $s$ and $t$.\n",
    "\n",
    "In one affine coupling layer, some dimensions (channels) remain unchanged. To make sure all the inputs have a chance to be altered, the model reverses the ordering in each layer so that different components are left unchanged. Following such an alternating pattern, the set of units which remain identical in one transformation layer are always modified in the next. \n",
    "\n",
    "This can be implemented with binary masks. First, we can extend the scale and neural networks to mappings form $\\mathbb{R}^D$ to $\\mathbb{R}^D$. Then taking a mask $\\mathbf{b} = (1,\\dots,1,0,\\dots,0)$ with $d$ ones, so that we have for the affine layer:\n",
    "\\begin{align*}\n",
    "\\mathbf{y} = \\mathbf{x} \\odot \\exp\\big((1-\\mathbf{b}) \\odot s(\\mathbf{b} \\odot \\mathbf{x})\\big) + (1-\\mathbf{b}) \\odot t(\\mathbf{b} \\odot \\mathbf{x}).\n",
    "\\end{align*}\n",
    "Note that we have\n",
    "\\begin{align*}\n",
    "\\ln \\left\\vert\\det(J(\\mathbf{x}))\\right\\vert = \\sum_{j=1}^{D} \\Big((1-\\mathbf{b})\\odot s(\\mathbf{b} \\odot \\mathbf{x})\\Big)_j,\n",
    "\\end{align*}\n",
    "and to invert the affine layer:\n",
    "\\begin{align*}\n",
    "\\mathbf{x} = \\left( \\mathbf{y} -(1-\\mathbf{b}) \\odot t(\\mathbf{b} \\odot \\mathbf{y})\\right)\\odot \\exp\\left( -(1-\\mathbf{b}) \\odot s(\\mathbf{b} \\odot \\mathbf{y})\\right)\n",
    "\\end{align*}\n",
    "Now we alternates the binary mask $\\mathbf{b}$ from one coupling layer to the other. \n",
    "\n",
    "Note, that the formula given in the paper is sligthly different:\n",
    "$$\\mathbf{y} = \\mathbf{b} \\odot \\mathbf{x} + (1 - \\mathbf{b}) \\odot \\Big(\\mathbf{x} \\odot \\exp\\big(s(\\mathbf{b} \\odot \\mathbf{x})\\big) + t(\\mathbf{b} \\odot \\mathbf{x})\\Big),$$\n",
    "but the 2 formulas give the same result!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NgMBEYdwbTS2"
   },
   "source": [
    "# Implementation of Real NVP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "loxXeNy9bTS3"
   },
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iFaqlfn3HcOV"
   },
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 87
    },
    "colab_type": "code",
    "id": "nFVr-ttwHcOX",
    "outputId": "cc585ee6-2740-487b-96fd-919507c45b5f"
   },
   "outputs": [],
   "source": [
    "from torch import distributions\n",
    "prior = distributions.MultivariateNormal(torch.zeros(2), torch.eye(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 87
    },
    "colab_type": "code",
    "id": "nFVr-ttwHcOX",
    "outputId": "cc585ee6-2740-487b-96fd-919507c45b5f"
   },
   "outputs": [],
   "source": [
    "# you can compute logprob and sample from your distribution:\n",
    "print(prior.log_prob(torch.Tensor([0,0])))\n",
    "print(prior.sample((3,)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GpyZ5AqObTS7"
   },
   "outputs": [],
   "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 dim(X)\n",
    "        # return x: a torch.Tensor of shape batchSize 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 dim(X)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "J6YPmM_3bTTE"
   },
   "outputs": [],
   "source": [
    "flow = RealNVP(nets, nett, masks, prior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that a flow is invertible g(f(x)) = x Hint: torch.allclose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vb80JOSSbTTG"
   },
   "outputs": [],
   "source": [
    "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?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6hOxCpXYbTTJ"
   },
   "source": [
    "# Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RmK9g7CIbTTK"
   },
   "outputs": [],
   "source": [
    "rcParams['figure.figsize'] = 10, 8"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qbhwTtHjpHC9"
   },
   "outputs": [],
   "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], x[:, 1], c='r')\n",
    "plt.title(r'$X = g(z)$')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Draw the estimated density:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xpoints = np.linspace(-1.5, 2.5, 500)\n",
    "ypoints = np.linspace(-1.0, 1.5, 500)\n",
    "(x1, x2,) = np.meshgrid(xpoints, ypoints)\n",
    "xgrid = np.concatenate((x1.reshape(-1, 1), x2.reshape(-1, 1)), axis=1).astype(np.float32)\n",
    "p = np.exp(flow.log_prob(torch.from_numpy(xgrid)).detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plt.imshow(\n",
    "    p.reshape(x1.shape), aspect=\"equal\", origin=\"lower\")\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Normalizing_flows_empty.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "dldiy",
   "language": "python",
   "name": "dldiy"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}