{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.distributions as dist\n",
    "from torch import nn\n",
    "from torchvision import datasets, transforms\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 512\n",
    "dl_train = torch.utils.data.DataLoader(\n",
    "    datasets.FashionMNIST('.', train=True, transform=transforms.ToTensor(), download=True),\n",
    "batch_size=batch_size, shuffle=True)\n",
    "\n",
    "dl_test = torch.utils.data.DataLoader(\n",
    "    datasets.FashionMNIST('.', train=False, transform=transforms.ToTensor(), download=True),\n",
    "batch_size=batch_size * 3, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5776 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6092 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5984 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "batch, _ = next(iter(dl_train))\n",
    "\n",
    "for i in range(3):\n",
    "    plt.imshow(1 - batch[i].view(28, 28).numpy(), cmap='Greys')\n",
    "    plt.show()\n",
    "    \n",
    "del batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def middle_size(size_in, size_out):\n",
    "    \"\"\"\n",
    "    Make a funnel with regard to the number of weights.\n",
    "    \"\"\"\n",
    "    return (size_in - size_out) // 2 + size_out\n",
    "\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self, latent_size):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.latent_size = latent_size \n",
    "        size = middle_size(784, self.latent_size)\n",
    "        \n",
    "        # Encode into gaussian paramaters. First part of vector is mu and second part log variance\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(784, size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(size, size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(size, latent_size * 2),\n",
    "        )\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(latent_size, size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(size, 784),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "        \n",
    "    def reparameterize(self, mu, log_var):\n",
    "        std = torch.exp(0.5 * log_var)\n",
    "        eps = torch.randn_like(std)  # unit gaussian\n",
    "        z = mu + eps * std \n",
    "        return z\n",
    "\n",
    "    def forward(self, x):\n",
    "        z_theta = self.encoder(x)\n",
    "        mu = z_theta[:, :self.latent_size]\n",
    "        log_var = z_theta[:, self.latent_size:]\n",
    "        z = self.reparameterize(mu, log_var)\n",
    "        return self.decoder(z), mu, log_var, z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(reconstruction_x, x, mu, log_var, z, mc_kl=True, beta=1.):\n",
    "    likelihood = F.binary_cross_entropy(reconstruction_x, x.view(-1, 784), reduction='sum')\n",
    "    \n",
    "    # is worse then the analyitical solution, but is better for comparing ELBO with Flows\n",
    "    if mc_kl:\n",
    "        log_prior = dist.Normal(0, 1).log_prob(z)\n",
    "        log_p_q = dist.Normal(mu, torch.exp(0.5 * log_var)).log_prob(z)\n",
    "        kl = (log_p_q - log_prior).sum()\n",
    "    else: \n",
    "        # see Appendix B from VAE paper:\n",
    "        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014\n",
    "        # https://arxiv.org/abs/1312.6114\n",
    "        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)\n",
    "        kl = -0.5 * torch.sum(1 + log_var - mu**2 - log_var.exp())\n",
    "\n",
    "    return (likelihood + beta * kl) / x.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3865041732788086"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def evaluate_test(device='cuda'):\n",
    "    avg_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for (x, _) in dl_test:\n",
    "            x = x.to(device).view(-1, 784)\n",
    "            reconstruction_x, mu, log_var, z = m(x)\n",
    "            avg_loss += loss_function(reconstruction_x, x, mu, log_var, z, 1)\n",
    "    return (avg_loss / len(dl_test.dataset)).item()\n",
    "  \n",
    "m = VAE(latent_size)\n",
    "m.cuda()\n",
    "evaluate_test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:   1\tAverage loss: 0.6467\tAverage test loss: 0.2048\tBeta: 1.000\n",
      "Epoch:   2\tAverage loss: 0.5570\tAverage test loss: 0.1945\tBeta: 1.000\n",
      "Epoch:   3\tAverage loss: 0.5292\tAverage test loss: 0.1865\tBeta: 1.000\n",
      "Epoch:   4\tAverage loss: 0.5123\tAverage test loss: 0.1829\tBeta: 1.000\n",
      "Epoch:   5\tAverage loss: 0.5071\tAverage test loss: 0.1807\tBeta: 1.000\n",
      "Epoch:   6\tAverage loss: 0.5020\tAverage test loss: 0.1797\tBeta: 1.000\n",
      "Epoch:   7\tAverage loss: 0.4995\tAverage test loss: 0.1791\tBeta: 1.000\n",
      "Epoch:   8\tAverage loss: 0.4980\tAverage test loss: 0.1784\tBeta: 1.000\n",
      "Epoch:   9\tAverage loss: 0.4963\tAverage test loss: 0.1785\tBeta: 1.000\n",
      "Epoch:  10\tAverage loss: 0.4959\tAverage test loss: 0.1778\tBeta: 1.000\n",
      "Epoch:  11\tAverage loss: 0.4949\tAverage test loss: 0.1776\tBeta: 1.000\n",
      "Epoch:  12\tAverage loss: 0.4937\tAverage test loss: 0.1780\tBeta: 1.000\n",
      "Epoch:  13\tAverage loss: 0.4939\tAverage test loss: 0.1784\tBeta: 1.000\n",
      "Epoch:  14\tAverage loss: 0.4935\tAverage test loss: 0.1787\tBeta: 1.000\n",
      "Epoch:  15\tAverage loss: 0.4928\tAverage test loss: 0.1764\tBeta: 1.000\n",
      "Epoch:  16\tAverage loss: 0.4909\tAverage test loss: 0.1767\tBeta: 1.000\n",
      "Epoch:  17\tAverage loss: 0.4902\tAverage test loss: 0.1758\tBeta: 1.000\n",
      "Epoch:  18\tAverage loss: 0.4896\tAverage test loss: 0.1757\tBeta: 1.000\n",
      "Epoch:  19\tAverage loss: 0.4904\tAverage test loss: 0.1791\tBeta: 1.000\n",
      "Epoch:  20\tAverage loss: 0.4893\tAverage test loss: 0.1760\tBeta: 1.000\n",
      "Epoch:  21\tAverage loss: 0.4885\tAverage test loss: 0.1759\tBeta: 1.000\n",
      "Epoch:  22\tAverage loss: 0.4881\tAverage test loss: 0.1755\tBeta: 1.000\n",
      "Epoch:  23\tAverage loss: 0.4881\tAverage test loss: 0.1764\tBeta: 1.000\n",
      "Epoch:  24\tAverage loss: 0.4883\tAverage test loss: 0.1750\tBeta: 1.000\n",
      "Epoch:  25\tAverage loss: 0.4873\tAverage test loss: 0.1753\tBeta: 1.000\n",
      "Epoch:  26\tAverage loss: 0.4882\tAverage test loss: 0.1749\tBeta: 1.000\n",
      "Epoch:  27\tAverage loss: 0.4890\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  28\tAverage loss: 0.4867\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  29\tAverage loss: 0.4867\tAverage test loss: 0.1748\tBeta: 1.000\n",
      "Epoch:  30\tAverage loss: 0.4868\tAverage test loss: 0.1753\tBeta: 1.000\n",
      "Epoch:  31\tAverage loss: 0.4867\tAverage test loss: 0.1755\tBeta: 1.000\n",
      "Epoch:  32\tAverage loss: 0.4871\tAverage test loss: 0.1746\tBeta: 1.000\n",
      "Epoch:  33\tAverage loss: 0.4863\tAverage test loss: 0.1749\tBeta: 1.000\n",
      "Epoch:  34\tAverage loss: 0.4861\tAverage test loss: 0.1746\tBeta: 1.000\n",
      "Epoch:  35\tAverage loss: 0.4865\tAverage test loss: 0.1763\tBeta: 1.000\n",
      "Epoch:  36\tAverage loss: 0.4864\tAverage test loss: 0.1745\tBeta: 1.000\n",
      "Epoch:  37\tAverage loss: 0.4870\tAverage test loss: 0.1750\tBeta: 1.000\n",
      "Epoch:  38\tAverage loss: 0.4861\tAverage test loss: 0.1748\tBeta: 1.000\n",
      "Epoch:  39\tAverage loss: 0.4859\tAverage test loss: 0.1750\tBeta: 1.000\n",
      "Epoch:  40\tAverage loss: 0.4861\tAverage test loss: 0.1745\tBeta: 1.000\n",
      "Epoch:  41\tAverage loss: 0.4860\tAverage test loss: 0.1771\tBeta: 1.000\n",
      "Epoch:  42\tAverage loss: 0.4857\tAverage test loss: 0.1747\tBeta: 1.000\n",
      "Epoch:  43\tAverage loss: 0.4855\tAverage test loss: 0.1779\tBeta: 1.000\n",
      "Epoch:  44\tAverage loss: 0.4868\tAverage test loss: 0.1748\tBeta: 1.000\n",
      "Epoch:  45\tAverage loss: 0.4858\tAverage test loss: 0.1746\tBeta: 1.000\n",
      "Epoch:  46\tAverage loss: 0.4856\tAverage test loss: 0.1745\tBeta: 1.000\n",
      "Epoch:  47\tAverage loss: 0.4859\tAverage test loss: 0.1747\tBeta: 1.000\n",
      "Epoch:  48\tAverage loss: 0.4855\tAverage test loss: 0.1743\tBeta: 1.000\n",
      "Epoch:  49\tAverage loss: 0.4853\tAverage test loss: 0.1743\tBeta: 1.000\n"
     ]
    }
   ],
   "source": [
    "latent_size = 32\n",
    "m = VAE(latent_size)\n",
    "m.cuda()\n",
    "\n",
    "vae_loss = []\n",
    "\n",
    "def log_training(batch_idx, x, loss):\n",
    "    if batch_idx % 400 == 0:\n",
    "        print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "            epoch, batch_idx * len(x), len(dl.dataset),\n",
    "            100. * batch_idx / len(dl),\n",
    "            loss.item() / len(x)))\n",
    "\n",
    "def train(epoch, m, device, optimizer, beta):\n",
    "    m.train()\n",
    "    train_loss = 0\n",
    "    for batch_idx, (x, _) in enumerate(dl_train):\n",
    "        x = x.to(device).view(-1, 784)\n",
    "        \n",
    "        optim.zero_grad()\n",
    "        reconstruction_x, mu, log_var, z = m(x)\n",
    "        loss = loss_function(reconstruction_x, x, mu, log_var, z, beta)\n",
    "        loss.backward()\n",
    "        \n",
    "        train_loss += loss.item()\n",
    "        optimizer.step()\n",
    "    \n",
    "    avg_loss = train_loss / len(dl_train.dataset)\n",
    "    print('Epoch: {:3d}'\n",
    "          '\\tAverage loss: {:.4f}'\n",
    "          '\\tAverage test loss: {:.4f}'\n",
    "          '\\tBeta: {:.3f}'.format(epoch, avg_loss, evaluate_test(), beta))\n",
    "    vae_loss.append(avg_loss)\n",
    "\n",
    "\n",
    "# optim = torch.optim.Adam(m.parameters(), 5e-4)\n",
    "\n",
    "# for epoch in range(1, 111):\n",
    "#     beta = min(1, 1 / 90 * epoch)\n",
    "#     train(epoch, m, 'cuda', optim, beta)\n",
    "    \n",
    "optim = torch.optim.Adam(m.parameters(), 1e-2)\n",
    "\n",
    "for epoch in range(1, 50):\n",
    "    beta = 1.\n",
    "    train(epoch, m, 'cuda', optim, beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6380 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6924 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6320 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6340 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6124 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6856 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6352 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6808 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 4772 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5536 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6308 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6288 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6272 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5560 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5788 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def show_samples(m, n=15, device='cuda'):\n",
    "    with torch.no_grad():\n",
    "        # take n samples\n",
    "        sample = torch.randn(n, latent_size).to(device)\n",
    "        sample = m.decoder(sample).cpu().view(-1, 28, 28)\n",
    "        \n",
    "    for i in range(sample.shape[0]):\n",
    "        plt.imshow(sample[i].view(28, 28), cmap='Greys')\n",
    "        plt.show()\n",
    "        \n",
    "show_samples(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic planar flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlanarFlow(nn.Module):\n",
    "    def __init__(self, size=1, init_sigma=0.01):\n",
    "        \"\"\"\n",
    "        shape u = (batch_size, z_size, 1)\n",
    "        shape w = (batch_size, 1, z_size)\n",
    "        shape b = (batch_size, 1, 1)\n",
    "        shape z = (batch_size, z_size).\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.u = nn.Parameter(torch.randn(size, 1).normal_(0, 0.01))\n",
    "        self.w = nn.Parameter(torch.randn(1, size).normal_(0, 0.01))\n",
    "        self.b = nn.Parameter(torch.zeros(1))\n",
    "    \n",
    "    @property\n",
    "    def normalized_u(self):\n",
    "        \"\"\"\n",
    "        Needed for invertibility condition.\n",
    "        \n",
    "        See Appendix A.1\n",
    "        Rezende et al. Variational Inference with Normalizing Flows\n",
    "        https://arxiv.org/pdf/1505.05770.pdf\n",
    "        \"\"\"\n",
    "        # softplus\n",
    "        def m(x):\n",
    "            return -1 + torch.log(1 + torch.exp(x))\n",
    "        wtu = self.w @ self.u\n",
    "        w_div_w2 = self.w.t() / torch.sum(self.w ** 2, dim=1, keepdim=True)\n",
    "        return self.u + (m(wtu) - wtu) * w_div_w2\n",
    "    \n",
    "    def psi(self, z):\n",
    "        \"\"\"\n",
    "        ψ(z) =h′(w^tz+b)w\n",
    "        \n",
    "        See eq(11)\n",
    "        Rezende et al. Variational Inference with Normalizing Flows\n",
    "        https://arxiv.org/pdf/1505.05770.pdf\n",
    "        \"\"\"\n",
    "        return self.h_prime(z @ self.w.t() + self.b) @ self.w\n",
    "    \n",
    "    def h(self, x):\n",
    "        return torch.tanh(x)\n",
    "        \n",
    "    def h_prime(self, z):\n",
    "        return 1 - torch.tanh(z)**2\n",
    "        \n",
    "    def forward(self, z):\n",
    "        if isinstance(z, tuple):\n",
    "            z, accumulating_ldj = z\n",
    "        else:\n",
    "            z, accumulating_ldj = z, 0\n",
    "        psi = self.psi(z)\n",
    "\n",
    "        u = self.normalized_u\n",
    "\n",
    "        # determinant of jacobian\n",
    "        det = (1 + psi @ u)\n",
    "\n",
    "        # log |det Jac|\n",
    "        ldj = torch.log(torch.abs(det) + 1e-6)\n",
    "        \n",
    "        wzb = z @ self.w.t() + self.b\n",
    " \n",
    "        fz = z + (u.t() * self.h(wzb))\n",
    "\n",
    "        return fz, ldj + accumulating_ldj\n",
    "\n",
    "\n",
    "\n",
    "class FlowVAE(VAE):\n",
    "    def __init__(self, latent_size, n_flows=20):\n",
    "        super().__init__(latent_size)\n",
    "                \n",
    "        self.flow = nn.Sequential(*[\n",
    "            PlanarFlow(latent_size) for _ in range(n_flows)\n",
    "        ])   \n",
    "            \n",
    "    def forward(self, x):\n",
    "        z_theta = self.encoder(x)\n",
    "        mu0 = z_theta[:, :self.latent_size]\n",
    "        log_var0 = z_theta[:, self.latent_size:]\n",
    "        z0 = self.reparameterize(mu0, log_var0)\n",
    "        zk, accumulating_ldj = self.flow(z0)\n",
    "        \n",
    "        return z0, zk, self.decoder(zk), accumulating_ldj, mu0, log_var0\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def det_loss(reconstruction_x, x, mu, log_var, z_0, z_k, ldj, beta=1.):\n",
    "    \"\"\"\n",
    "    :param z_mu: mean of z_0\n",
    "    :param z_var: variance of z_0\n",
    "    :param z_0: first stochastic latent variable\n",
    "    :param z_k: last stochastic latent variable\n",
    "    :param ldj: log det jacobian\n",
    "    \"\"\"\n",
    "\n",
    "    batch_size = x.size(0)\n",
    "\n",
    "    # - N E_q0 [ ln p(x|z_k) ]\n",
    "    likelihood = F.binary_cross_entropy(reconstruction_x, x, reduction='sum')\n",
    "\n",
    "    # ln p(z_k)  (not averaged)\n",
    "    log_p_zk = dist.Normal(0, 1).log_prob(z_k)\n",
    "    # ln q(z_0)  (not averaged)\n",
    "    log_q_z0 = dist.Normal(mu, torch.exp(0.5 * log_var)).log_prob(z_0)\n",
    "    \n",
    "    return (log_q_z0.sum() + beta * (likelihood - log_p_zk.sum()) - ldj.sum()) / batch_size\n",
    "    \n",
    "#     # N E_q0[ ln q(z_0) - ln p(z_k) ]\n",
    "#     summed_logs = torch.sum(log_q_z0 - log_p_zk)\n",
    "\n",
    "#     # sum over batches\n",
    "#     summed_ldj = torch.sum(ldj)\n",
    "\n",
    "#     # ldj = N E_q_z0[\\sum_k log |det dz_k/dz_k-1| ]\n",
    "#     kl = (summed_logs - summed_ldj)\n",
    "#     loss = likelihood + beta * kl\n",
    "\n",
    "#     return loss / x.shape[0]\n",
    "\n",
    "def evaluate_test(device='cuda'):\n",
    "    avg_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for (x, _) in dl_test:\n",
    "            x = x.to(device).view(-1, 784)\n",
    "            z0, zk, reconstruction_x, ldj, mu0, log_var0 = m(x)\n",
    "\n",
    "            avg_loss += det_loss(x=x, \n",
    "                            reconstruction_x=reconstruction_x, \n",
    "                            mu=mu0,\n",
    "                            log_var=log_var0,\n",
    "                            ldj=ldj,\n",
    "                            z_0=z0,\n",
    "                            z_k=zk,\n",
    "                            beta=beta\n",
    "                            )\n",
    "    return (avg_loss / len(dl_test.dataset)).item()\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:   1\tAverage loss: 3.0833\tAverage test loss: 0.1994\tBeta: 1.000\n",
      "Epoch:   2\tAverage loss: 0.5354\tAverage test loss: 0.1869\tBeta: 1.000\n",
      "Epoch:   3\tAverage loss: 0.5139\tAverage test loss: 0.1829\tBeta: 1.000\n",
      "Epoch:   4\tAverage loss: 0.5071\tAverage test loss: 0.1813\tBeta: 1.000\n",
      "Epoch:   5\tAverage loss: 0.5028\tAverage test loss: 0.1803\tBeta: 1.000\n",
      "Epoch:   6\tAverage loss: 0.5001\tAverage test loss: 0.1792\tBeta: 1.000\n",
      "Epoch:   7\tAverage loss: 0.4980\tAverage test loss: 0.1787\tBeta: 1.000\n",
      "Epoch:   8\tAverage loss: 0.4967\tAverage test loss: 0.1783\tBeta: 1.000\n",
      "Epoch:   9\tAverage loss: 0.4948\tAverage test loss: 0.1781\tBeta: 1.000\n",
      "Epoch:  10\tAverage loss: 0.4949\tAverage test loss: 0.1773\tBeta: 1.000\n",
      "Epoch:  11\tAverage loss: 0.4922\tAverage test loss: 0.1765\tBeta: 1.000\n",
      "Epoch:  12\tAverage loss: 0.4918\tAverage test loss: 0.1761\tBeta: 1.000\n",
      "Epoch:  13\tAverage loss: 0.4906\tAverage test loss: 0.1764\tBeta: 1.000\n",
      "Epoch:  14\tAverage loss: 0.4902\tAverage test loss: 0.1762\tBeta: 1.000\n",
      "Epoch:  15\tAverage loss: 0.4898\tAverage test loss: 0.1760\tBeta: 1.000\n",
      "Epoch:  16\tAverage loss: 0.4900\tAverage test loss: 0.1783\tBeta: 1.000\n",
      "Epoch:  17\tAverage loss: 0.4888\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  18\tAverage loss: 0.4879\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  19\tAverage loss: 0.4880\tAverage test loss: 0.1757\tBeta: 1.000\n",
      "Epoch:  20\tAverage loss: 0.4882\tAverage test loss: 0.1749\tBeta: 1.000\n",
      "Epoch:  21\tAverage loss: 0.4874\tAverage test loss: 0.1752\tBeta: 1.000\n",
      "Epoch:  22\tAverage loss: 0.4869\tAverage test loss: 0.1749\tBeta: 1.000\n",
      "Epoch:  23\tAverage loss: 0.4868\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  24\tAverage loss: 0.4868\tAverage test loss: 0.1747\tBeta: 1.000\n",
      "Epoch:  25\tAverage loss: 0.4864\tAverage test loss: 0.1758\tBeta: 1.000\n",
      "Epoch:  26\tAverage loss: 0.4861\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  27\tAverage loss: 0.4860\tAverage test loss: 0.1749\tBeta: 1.000\n",
      "Epoch:  28\tAverage loss: 0.4860\tAverage test loss: 0.1742\tBeta: 1.000\n",
      "Epoch:  29\tAverage loss: 0.4852\tAverage test loss: 0.1757\tBeta: 1.000\n",
      "Epoch:  30\tAverage loss: 0.4847\tAverage test loss: 0.1768\tBeta: 1.000\n",
      "Epoch:  31\tAverage loss: 0.4854\tAverage test loss: 0.1748\tBeta: 1.000\n",
      "Epoch:  32\tAverage loss: 0.4852\tAverage test loss: 0.1747\tBeta: 1.000\n",
      "Epoch:  33\tAverage loss: 0.4842\tAverage test loss: 0.1737\tBeta: 1.000\n",
      "Epoch:  34\tAverage loss: 0.4842\tAverage test loss: 0.1735\tBeta: 1.000\n",
      "Epoch:  35\tAverage loss: 0.4836\tAverage test loss: 0.1755\tBeta: 1.000\n",
      "Epoch:  36\tAverage loss: 0.4840\tAverage test loss: 0.1738\tBeta: 1.000\n",
      "Epoch:  37\tAverage loss: 0.4835\tAverage test loss: 0.1739\tBeta: 1.000\n",
      "Epoch:  38\tAverage loss: 0.4841\tAverage test loss: 0.1748\tBeta: 1.000\n",
      "Epoch:  39\tAverage loss: 0.4827\tAverage test loss: 0.1750\tBeta: 1.000\n",
      "Epoch:  40\tAverage loss: 0.4837\tAverage test loss: 0.1741\tBeta: 1.000\n",
      "Epoch:  41\tAverage loss: 0.4826\tAverage test loss: 0.1734\tBeta: 1.000\n",
      "Epoch:  42\tAverage loss: 0.4832\tAverage test loss: 0.1733\tBeta: 1.000\n",
      "Epoch:  43\tAverage loss: 0.4828\tAverage test loss: 0.1742\tBeta: 1.000\n",
      "Epoch:  44\tAverage loss: 0.4833\tAverage test loss: 0.1733\tBeta: 1.000\n",
      "Epoch:  45\tAverage loss: 0.4833\tAverage test loss: 0.1742\tBeta: 1.000\n",
      "Epoch:  46\tAverage loss: 0.4827\tAverage test loss: 0.1740\tBeta: 1.000\n",
      "Epoch:  47\tAverage loss: 0.4818\tAverage test loss: 0.1731\tBeta: 1.000\n",
      "Epoch:  48\tAverage loss: 0.4826\tAverage test loss: 0.1731\tBeta: 1.000\n",
      "Epoch:  49\tAverage loss: 0.4825\tAverage test loss: 0.1733\tBeta: 1.000\n"
     ]
    }
   ],
   "source": [
    "m = FlowVAE(latent_size, n_flows=32)\n",
    "m.cuda()\n",
    "\n",
    "flow_loss = []\n",
    "\n",
    "def train(epoch, m, device, optimizer, beta):\n",
    "    m.train()\n",
    "    train_loss = 0\n",
    "    for batch_idx, (x, _) in enumerate(dl_train):\n",
    "        x = x.to(device).view(-1, 784)\n",
    "        \n",
    "        optim.zero_grad()\n",
    "        z0, zk, reconstruction_x, ldj, mu0, log_var0 = m(x)\n",
    "    \n",
    "        loss = det_loss(x=x, \n",
    "                        reconstruction_x=reconstruction_x, \n",
    "                        mu=mu0,\n",
    "                        log_var=log_var0,\n",
    "                        ldj=ldj,\n",
    "                        z_0=z0,\n",
    "                        z_k=zk,\n",
    "                        beta=beta\n",
    "                        )\n",
    "        loss.backward()\n",
    "        \n",
    "        train_loss += loss.item()\n",
    "        optimizer.step()\n",
    "    \n",
    "    avg_loss = train_loss / len(dl_train.dataset)\n",
    "    print('Epoch: {:3d}'\n",
    "          '\\tAverage loss: {:.4f}'\n",
    "          '\\tAverage test loss: {:.4f}'\n",
    "          '\\tBeta: {:.3f}'.format(epoch, avg_loss, evaluate_test(), beta))\n",
    "    vae_loss.append(avg_loss)\n",
    "\n",
    "\n",
    "optim = torch.optim.Adam(m.parameters(), 1e-2)\n",
    "\n",
    "for epoch in range(1, 50):\n",
    "#     beta = min(1, 1 / 90 * epoch)\n",
    "    beta = 1.\n",
    "    train(epoch, m, 'cuda', optim, beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5932 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6644 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6808 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6320 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5640 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6412 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6448 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6020 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5740 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5800 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5404 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5768 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6128 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6432 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6116 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_samples(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38576 characters]\n",
      "text/plain": [
       "<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=101x128 at 0x7FE421AD15C0>"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img = Image.open('/home/ritchie46/Downloads/Donald_Trump_official_portrait.jpg')\n",
    "img.thumbnail((128, 128))\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = img.convert('1')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fe431105588>"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12240 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = 1 - np.array(img)\n",
    "plt.imshow(img, cmap=\"Greys\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "img.shape\n",
    "img = torch.tensor(img, dtype=torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_z(z):\n",
    "    a = z.numpy().reshape(img.shape)\n",
    "    plt.imshow(a, cmap=\"Greys\")\n",
    "    plt.show()\n",
    "    \n",
    "size = img.shape[0] * img.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48284 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48416 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "class Flow(nn.Module):\n",
    "    def __init__(self, n_flows=10):\n",
    "        super().__init__()\n",
    "        self.flow =  nn.Sequential(*[\n",
    "            PlanarFlow(size) for _ in range(n_flows)\n",
    "        ])\n",
    "        self.mu = nn.Parameter(torch.randn(size,).normal_(0, 0.01))\n",
    "        self.log_var = nn.Parameter(torch.randn(size,).normal_(0, 0.01))\n",
    "        \n",
    "    def forward(self):\n",
    "        std = torch.exp(0.5 * self.log_var)\n",
    "        eps = torch.randn_like(std)  # unit gaussian\n",
    "        z0 = self.mu + eps * std \n",
    "        zk, ldj = self.flow(z0)\n",
    "        return z0, zk, ldj, self.mu, self.log_var\n",
    "        \n",
    "\n",
    "flow = Flow()\n",
    "z0, zk, ldj, mu, log_var = flow()\n",
    "\n",
    "show_z(z0.data)\n",
    "show_z(zk.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "flow = Flow()\n",
    "\n",
    "def train_flow(flow, epochs=250):\n",
    "    optim = torch.optim.Adam(flow.parameters(), lr=1e-3)\n",
    "    for _ in range(epochs):\n",
    "\n",
    "        z0, zk, ldj, mu, log_var = flow()\n",
    "\n",
    "        loss = det_loss(reconstruction_x=torch.sigmoid(zk),\n",
    "                       x=img,\n",
    "                       mu=mu,\n",
    "                       log_var=log_var,\n",
    "                       z_0=z0,\n",
    "                       z_k=zk,\n",
    "                       ldj=ldj)\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "        optim.zero_grad()  \n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:13: UserWarning: Using a target size (torch.Size([128, 101])) that is different to the input size (torch.Size([1, 12928])) is deprecated. Please ensure they have the same size.\n",
      "  del sys.path[0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N flows: 5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49656 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N flows: 20\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49576 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N flows: 32\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49592 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N flows: 64\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49636 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for n in [5, 20, 32, 64]:\n",
    "    flow = Flow(n)\n",
    "    train_flow(flow)\n",
    "    z0, zk, ldj, mu, log_var = flow()\n",
    "    print('N flows:', n)\n",
    "    show_z(torch.sigmoid(zk).data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amortized planar flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/riannevdberg/sylvester-flows\n",
    "\n",
    "class Planar(nn.Module):\n",
    "    \"\"\"\n",
    "    PyTorch implementation of planar flows as presented in \"Variational Inference with Normalizing Flows\"\n",
    "    by Danilo Jimenez Rezende, Shakir Mohamed. Model assumes amortized flow parameters.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "\n",
    "        super(Planar, self).__init__()\n",
    "\n",
    "        self.h = nn.Tanh()\n",
    "        self.softplus = nn.Softplus()\n",
    "\n",
    "    def der_h(self, x):\n",
    "        \"\"\" Derivative of tanh \"\"\"\n",
    "\n",
    "        return 1 - self.h(x) ** 2\n",
    "\n",
    "    def forward(self, zk, u, w, b):\n",
    "        \"\"\"\n",
    "        Forward pass. Assumes amortized u, w and b. Conditions on diagonals of u and w for invertibility\n",
    "        will be be satisfied inside this function. Computes the following transformation:\n",
    "        z' = z + u h( w^T z + b)\n",
    "        or actually\n",
    "        z'^T = z^T + h(z^T w + b)u^T\n",
    "        Assumes the following input shapes:\n",
    "        shape u = (batch_size, z_size, 1)\n",
    "        shape w = (batch_size, 1, z_size)\n",
    "        shape b = (batch_size, 1, 1)\n",
    "        shape z = (batch_size, z_size).\n",
    "        \"\"\"\n",
    "\n",
    "        zk = zk.unsqueeze(2)\n",
    "\n",
    "        # reparameterize u such that the flow becomes invertible (see appendix paper)\n",
    "        uw = torch.bmm(w, u)\n",
    "        m_uw = -1. + self.softplus(uw)\n",
    "        w_norm_sq = torch.sum(w ** 2, dim=2, keepdim=True)\n",
    "        u_hat = u + ((m_uw - uw) * w.transpose(2, 1) / w_norm_sq)\n",
    "\n",
    "        # compute flow with u_hat\n",
    "        wzb = torch.bmm(w, zk) + b\n",
    "        z = zk + u_hat * self.h(wzb)\n",
    "        z = z.squeeze(2)\n",
    "\n",
    "        # compute logdetJ\n",
    "        psi = w * self.der_h(wzb)\n",
    "        log_det_jacobian = torch.log(torch.abs(1 + torch.bmm(psi, u_hat)))\n",
    "        log_det_jacobian = log_det_jacobian.squeeze(2).squeeze(1)\n",
    "\n",
    "        return z, log_det_jacobian\n",
    "\n",
    "class PlanarVAE(VAE):\n",
    "    \"\"\"\n",
    "    Variational auto-encoder with planar flows in the encoder.\n",
    "    \"\"\"\n",
    "    def __init__(self, latent_size, n_flows):\n",
    "        super(PlanarVAE, self).__init__(latent_size)\n",
    "\n",
    "        # Initialize log-det-jacobian to zero\n",
    "        self.log_det_j = 0.\n",
    "\n",
    "        # Flow parameters\n",
    "        flow = Planar\n",
    "        self.n_flows = n_flows\n",
    "\n",
    "        # Amortized flow parameters\n",
    "        self.amor_u = nn.Linear(latent_size * 2, n_flows * latent_size)\n",
    "        self.amor_w = nn.Linear(latent_size * 2, n_flows * latent_size)\n",
    "        self.amor_b = nn.Linear(latent_size * 2, n_flows)\n",
    "\n",
    "        # Normalizing flow layers\n",
    "        for k in range(n_flows):\n",
    "            flow_k = flow()\n",
    "            self.add_module('flow_' + str(k), flow_k)\n",
    "\n",
    "    def encode(self, x):\n",
    "        \"\"\"\n",
    "        Encoder that ouputs parameters for base distribution of z and flow parameters.\n",
    "        \"\"\"\n",
    "\n",
    "        batch_size = x.size(0)\n",
    "\n",
    "        h = self.q_z_nn(x)\n",
    "        h = h.view(-1, self.q_z_nn_output_dim)\n",
    "        mean_z = self.q_z_mean(h)\n",
    "        var_z = self.q_z_var(h)\n",
    "\n",
    "        # return amortized u an w for all flows\n",
    "        u = self.amor_u(h).view(batch_size, self.num_flows, self.z_size, 1)\n",
    "        w = self.amor_w(h).view(batch_size, self.num_flows, 1, self.z_size)\n",
    "        b = self.amor_b(h).view(batch_size, self.num_flows, 1, 1)\n",
    "\n",
    "        return mean_z, var_z, u, w, b\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward pass with planar flows for the transformation z_0 -> z_1 -> ... -> z_k.\n",
    "        Log determinant is computed as log_det_j = N E_q_z0[\\sum_k log |det dz_k/dz_k-1| ].\n",
    "        \"\"\"\n",
    "        batch_size = x.shape[0]\n",
    "        z_theta = self.encoder(x)\n",
    "        mu0 = z_theta[:, :self.latent_size]\n",
    "        log_var0 = z_theta[:, self.latent_size:]\n",
    "        \n",
    "        u = self.amor_u(z_theta).view(batch_size, self.n_flows, self.latent_size, 1)\n",
    "        w = self.amor_w(z_theta).view(batch_size, self.n_flows, 1, self.latent_size)\n",
    "        b = self.amor_b(z_theta).view(batch_size, self.n_flows, 1, 1)\n",
    "            \n",
    "        self.log_det_j = 0.\n",
    "        z = [self.reparameterize(mu0, log_var0)]\n",
    "        \n",
    "        # Normalizing flows\n",
    "        for k in range(self.n_flows):\n",
    "            flow_k = getattr(self, 'flow_' + str(k))\n",
    "            z_k, log_det_jacobian = flow_k(z[k], u[:, k, :, :], w[:, k, :, :], b[:, k, :, :])\n",
    "            z.append(z_k)\n",
    "            self.log_det_j += log_det_jacobian\n",
    "            \n",
    "        z0 = z[0]\n",
    "        zk = z[-1]\n",
    "        \n",
    "        return z0, zk, self.decoder(zk), self.log_det_j, mu0, log_var0\n",
    "\n",
    "m = PlanarVAE(20, 10)\n",
    "m(next(iter(dl_train))[0].view(-1, 784));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:   1\tAverage loss: 0.8979\tAverage test loss: 0.2050\tBeta: 1.000\n",
      "Epoch:   2\tAverage loss: 0.5478\tAverage test loss: 0.1903\tBeta: 1.000\n",
      "Epoch:   3\tAverage loss: 0.5368\tAverage test loss: 0.1878\tBeta: 1.000\n",
      "Epoch:   4\tAverage loss: 0.5143\tAverage test loss: 0.1835\tBeta: 1.000\n",
      "Epoch:   5\tAverage loss: 0.5078\tAverage test loss: 0.1828\tBeta: 1.000\n",
      "Epoch:   6\tAverage loss: 0.5040\tAverage test loss: 0.1798\tBeta: 1.000\n",
      "Epoch:   7\tAverage loss: 0.5005\tAverage test loss: 0.1790\tBeta: 1.000\n",
      "Epoch:   8\tAverage loss: 0.5002\tAverage test loss: 0.1793\tBeta: 1.000\n",
      "Epoch:   9\tAverage loss: 0.4978\tAverage test loss: 0.1775\tBeta: 1.000\n",
      "Epoch:  10\tAverage loss: 0.4955\tAverage test loss: 0.1787\tBeta: 1.000\n",
      "Epoch:  11\tAverage loss: 0.4955\tAverage test loss: 0.1772\tBeta: 1.000\n",
      "Epoch:  12\tAverage loss: 0.4942\tAverage test loss: 0.1812\tBeta: 1.000\n",
      "Epoch:  13\tAverage loss: 0.4935\tAverage test loss: 0.1772\tBeta: 1.000\n",
      "Epoch:  14\tAverage loss: 0.4923\tAverage test loss: 0.1766\tBeta: 1.000\n",
      "Epoch:  15\tAverage loss: 0.4933\tAverage test loss: 0.1766\tBeta: 1.000\n",
      "Epoch:  16\tAverage loss: 0.4916\tAverage test loss: 0.1761\tBeta: 1.000\n",
      "Epoch:  17\tAverage loss: 0.5045\tAverage test loss: 0.1948\tBeta: 1.000\n",
      "Epoch:  18\tAverage loss: 0.5037\tAverage test loss: 0.1804\tBeta: 1.000\n",
      "Epoch:  19\tAverage loss: 0.4925\tAverage test loss: 0.1759\tBeta: 1.000\n",
      "Epoch:  20\tAverage loss: 0.4901\tAverage test loss: 0.1756\tBeta: 1.000\n",
      "Epoch:  21\tAverage loss: 0.4897\tAverage test loss: 0.1761\tBeta: 1.000\n",
      "Epoch:  22\tAverage loss: 0.4887\tAverage test loss: 0.1763\tBeta: 1.000\n",
      "Epoch:  23\tAverage loss: 0.4892\tAverage test loss: 0.1750\tBeta: 1.000\n",
      "Epoch:  24\tAverage loss: 0.4878\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  25\tAverage loss: 0.4884\tAverage test loss: 0.1755\tBeta: 1.000\n",
      "Epoch:  26\tAverage loss: 0.4887\tAverage test loss: 0.1760\tBeta: 1.000\n",
      "Epoch:  27\tAverage loss: 0.4886\tAverage test loss: 0.1751\tBeta: 1.000\n",
      "Epoch:  28\tAverage loss: 0.4876\tAverage test loss: 0.1758\tBeta: 1.000\n",
      "Epoch:  29\tAverage loss: 0.4872\tAverage test loss: 0.1751\tBeta: 1.000\n",
      "Epoch:  30\tAverage loss: 0.4870\tAverage test loss: 0.1745\tBeta: 1.000\n",
      "Epoch:  31\tAverage loss: 0.4873\tAverage test loss: 0.1756\tBeta: 1.000\n",
      "Epoch:  32\tAverage loss: 0.4887\tAverage test loss: 0.1745\tBeta: 1.000\n",
      "Epoch:  33\tAverage loss: 0.4862\tAverage test loss: 0.1755\tBeta: 1.000\n",
      "Epoch:  34\tAverage loss: 0.4867\tAverage test loss: 0.1757\tBeta: 1.000\n",
      "Epoch:  35\tAverage loss: 0.4868\tAverage test loss: 0.1753\tBeta: 1.000\n",
      "Epoch:  36\tAverage loss: 0.4867\tAverage test loss: 0.1759\tBeta: 1.000\n",
      "Epoch:  37\tAverage loss: 0.4867\tAverage test loss: 0.1747\tBeta: 1.000\n",
      "Epoch:  38\tAverage loss: 0.4872\tAverage test loss: 0.1766\tBeta: 1.000\n",
      "Epoch:  39\tAverage loss: 0.4865\tAverage test loss: 0.1760\tBeta: 1.000\n",
      "Epoch:  40\tAverage loss: 0.4865\tAverage test loss: 0.1751\tBeta: 1.000\n",
      "Epoch:  41\tAverage loss: 0.4858\tAverage test loss: 0.1761\tBeta: 1.000\n",
      "Epoch:  42\tAverage loss: 0.4868\tAverage test loss: 0.1754\tBeta: 1.000\n",
      "Epoch:  43\tAverage loss: 0.4871\tAverage test loss: 0.1771\tBeta: 1.000\n",
      "Epoch:  44\tAverage loss: 0.4866\tAverage test loss: 0.1749\tBeta: 1.000\n",
      "Epoch:  45\tAverage loss: 0.4858\tAverage test loss: 0.1748\tBeta: 1.000\n",
      "Epoch:  46\tAverage loss: 0.4856\tAverage test loss: 0.1747\tBeta: 1.000\n",
      "Epoch:  47\tAverage loss: 0.4855\tAverage test loss: 0.1750\tBeta: 1.000\n",
      "Epoch:  48\tAverage loss: 0.4850\tAverage test loss: 0.1740\tBeta: 1.000\n",
      "Epoch:  49\tAverage loss: 0.4854\tAverage test loss: 0.1747\tBeta: 1.000\n"
     ]
    }
   ],
   "source": [
    "m = PlanarVAE(latent_size, 16)\n",
    "m.cuda()\n",
    "optim = torch.optim.Adam(m.parameters(), 1e-2)\n",
    "\n",
    "for epoch in range(1, 50):\n",
    "    beta = 1.\n",
    "    train(epoch, m, 'cuda', optim, beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 391,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6016 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6804 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6628 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6560 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6304 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 7096 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6852 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6912 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6292 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6288 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6764 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6668 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6528 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6420 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6448 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_samples(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
