day2/notebooks/solutions_VAE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/PGM-Lab/2023-probai-private/blob/main/Day2-Evening/notebooks/solutions_VAE.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wtZLhDvtfCrP"
   },
   "source": [
    "# Adapt the variational auto encoder\n",
    "\n",
    "Below you will find an implementation of a VAE for the MNIST data. To allow for faster learning time, we only consider the digits 0,1, and 2 and only the first 100 samples of those digits.\n",
    "\n",
    "In this exercise, you should familiarize yourself with the implementation below and experiment with the structure of the VAE specification in order to emphasize digit separation in the latent space and the generation of images when sampling from the latent space.\n",
    "\n",
    "Part of the implementation is based on code from the official Pyro examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JsT1zW1sfCrT",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#!pip install -q --upgrade pyro-ppl torch torchvision \n",
    "!pip install -q pyro-ppl torch torchvision matplotlib numpy scipy seaborn pandas\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision.datasets as datasets\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as transforms\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "from pyro.infer import SVI, Trace_ELBO\n",
    "from pyro.optim import Adam\n",
    "import datetime\n",
    "import os\n",
    "import matplotlib.gridspec as gridspec\n",
    "from matplotlib import pyplot\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from scipy.stats import norm\n",
    "import copy\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "\n",
    "import ssl\n",
    "ssl._create_default_https_context = ssl._create_unverified_context\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4cHlQ9xEfCrU"
   },
   "source": [
    "### Get the MNIST data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y2Y0vokLfCrU",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# For reproducibility\n",
    "torch.manual_seed(1234)\n",
    "np.random.seed(1234)\n",
    "\n",
    "# By default, we only select the digits 5, 6, and 7 and only 100 examples from each of these\n",
    "# digit classes\n",
    "def get_data_sample(digits = [5,6,7], number_of_images_from_each_class = 100):\n",
    "\n",
    "    _data = datasets.MNIST('./data', train=True, download=True)\n",
    "\n",
    "    selector = np.array([], dtype=int)\n",
    "    for i in digits:\n",
    "        selector = np.concatenate((selector, np.random.choice(np.where(_data.targets == i)[0], \n",
    "                                                              size=number_of_images_from_each_class, replace=False)))\n",
    "\n",
    "    _data.data = _data.data[selector, :, :]\n",
    "    _data.targets = _data.targets[selector]\n",
    "\n",
    "    # Binarize the data\n",
    "    _data.data[_data.data<128] = 0\n",
    "    _data.data[_data.data>=128] = 1\n",
    "\n",
    "    _data.data = _data.data.type(torch.float)\n",
    "    _data.data = _data.data.reshape(number_of_images_from_each_class*len(digits),-1)\n",
    "    \n",
    "    return _data, digits, number_of_images_from_each_class\n",
    "\n",
    "data, digits, number_of_images_from_each_class = get_data_sample(number_of_images_from_each_class=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 248
    },
    "id": "FXyb663MfCrU",
    "outputId": "48e268cd-a125-485b-a409-e5431989bf22",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def display_image(x):\n",
    "    plt.figure(figsize=(3,3))\n",
    "    plt.imshow(x.reshape((28, 28)), cmap=\"gray\")\n",
    "\n",
    "toy_image = data.data[0,:]\n",
    "display_image(toy_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3iPRoMFgfCrZ"
   },
   "source": [
    "# Defining a VAE in Pyro\n",
    "\n",
    "In the following implementation, the econder and decoder networks are provided as paramters to the VAE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6_jKIv68fCrZ",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class VAE(nn.Module):\n",
    "    def __init__(self, z_dim, encoder, decoder):\n",
    "        super(VAE, self).__init__()\n",
    "        # set the encoder and decoder networks\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder        \n",
    "        self.z_dim = z_dim\n",
    "\n",
    "    # define the model p(x|z)p(z)\n",
    "    def model(self, x):\n",
    "        # register PyTorch module `decoder` with Pyro\n",
    "        pyro.module(\"decoder\", self.decoder)\n",
    "        with pyro.plate(\"data\", x.shape[0]):\n",
    "            # setup hyperparameters for prior p(z)\n",
    "            z_loc = x.new_zeros(torch.Size((x.shape[0], self.z_dim)))\n",
    "            z_scale = x.new_ones(torch.Size((x.shape[0], self.z_dim)))\n",
    "            # sample from prior (value will be sampled by guide when computing the ELBO)\n",
    "            z = pyro.sample(\"latent\", dist.Normal(z_loc, z_scale).to_event(1))\n",
    "            # decode the latent code z\n",
    "            loc_img = self.decoder.forward(z)\n",
    "            # score against actual images\n",
    "            pyro.sample(\"obs\", dist.Bernoulli(logits=loc_img).to_event(1), obs=x.reshape(-1, 784))\n",
    "\n",
    "    # define the guide (i.e. variational distribution) q(z|x)\n",
    "    def guide(self, x):\n",
    "        # register PyTorch module `encoder` with Pyro\n",
    "        pyro.module(\"encoder\", self.encoder)\n",
    "        with pyro.plate(\"data\", x.shape[0]):\n",
    "            # use the encoder to get the parameters used to define q(z|x)\n",
    "            z_loc, z_scale = self.encoder.forward(x)\n",
    "            # sample the latent code z\n",
    "            pyro.sample(\"latent\", dist.Normal(z_loc, z_scale).to_event(1))\n",
    "    \n",
    "    def sample_images(self, z, num_images = 5):\n",
    "        if type(z) is not torch.Tensor: \n",
    "            z=torch.tensor(z, dtype=torch.float32)\n",
    "   \n",
    "        # Pass Z through the decoder network\n",
    "        logits = self.decoder.forward(z)\n",
    "\n",
    "        # Bernoulli distribution over images\n",
    "        distribution_over_images = dist.Bernoulli(logits=logits)\n",
    "\n",
    "        f, ax = plt.subplots(1, num_images+1, figsize=(10, 10), sharex=True)\n",
    "        plt.axis('off')\n",
    "        ax[0].imshow(distribution_over_images.mean.detach().numpy().reshape((28, 28)), cmap=\"gray\")\n",
    "        ax[0].set_xlabel('Mean of Distribution')\n",
    "        for i in range(1, num_images+1):\n",
    "          ax[i].imshow(distribution_over_images.sample().detach().numpy().reshape((28, 28)), cmap=\"gray\")\n",
    "          ax[i].set_xlabel('Sample '+str(i))\n",
    "        plt.show()        \n",
    "        \n",
    "    def reconstruct_images(self, num_images=5):\n",
    "        # Sample num_images randomly and (try) to reconstruct them.\n",
    "        select_images_idx = np.random.choice(data.data.shape[0], size=num_images, replace=False)\n",
    "        f, ax = plt.subplots(1, num_images, figsize=(10, 10))\n",
    "        plt.tight_layout()\n",
    "        # Display sampled images\n",
    "        for idx, i in enumerate(select_images_idx):\n",
    "            toy_image = data.data[i,:]\n",
    "            ax[idx].imshow(toy_image.reshape(28,28), cmap=\"gray\")    \n",
    "        plt.show()\n",
    "        # Display reconstructed images\n",
    "        f, ax = plt.subplots(1, num_images, figsize=(10, 10))\n",
    "        plt.tight_layout()\n",
    "        for idx, i in enumerate(select_images_idx):\n",
    "            toy_image = data.data[i,:]\n",
    "            z_loc, _ = self.encoder(toy_image)\n",
    "            logits_img = self.decoder(z_loc) \n",
    "            img_mean = dist.Bernoulli(logits=logits_img).mean\n",
    "            ax[idx].imshow(img_mean.detach().numpy().reshape(28,28), cmap=\"gray\")\n",
    "        plt.show()   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p6QS-heeBE2-"
   },
   "source": [
    "# Defining the Encoder\n",
    "\n",
    "Define a simple encoder with a single hidden layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qnNSDbvofCrY",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, z_dim, hidden_dim):\n",
    "        super(Encoder, self).__init__()\n",
    "        # setup the three linear transformations used\n",
    "        self.fc1 = nn.Linear(784, hidden_dim)\n",
    "        self.fc21 = nn.Linear(hidden_dim, z_dim)\n",
    "        self.fc22 = nn.Linear(hidden_dim, z_dim)\n",
    "        # setup the non-linearities\n",
    "        self.non_linearity = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # define the forward computation on the image x\n",
    "        # first shape the mini-batch to have pixels in the rightmost dimension\n",
    "        x = x.reshape(-1, 784)\n",
    "        # then compute the hidden units\n",
    "        hidden = self.non_linearity(self.fc1(x))\n",
    "        # then return a mean vector and a (positive) square root covariance\n",
    "        # each of size batch_size x z_dim\n",
    "        z_loc = self.fc21(hidden)\n",
    "        z_scale =nn.functional.softplus(self.fc22(hidden))\n",
    "        return z_loc, z_scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tNxWgQVtfCrW"
   },
   "source": [
    "# Defining the Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZXaj5hQJ_-Zg"
   },
   "source": [
    "Define a linear decoder (i.e. a linear mapping) between the latent representation $z$ and the output space $x$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tFe4aDdgfCrW",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class LinearDecoder(nn.Module):\n",
    "    def __init__(self, z_dim):\n",
    "        super(LinearDecoder, self).__init__()\n",
    "        # setup the two linear transformations used\n",
    "        self.fc1 = nn.Linear(z_dim, 784)\n",
    "\n",
    "    def forward(self, z):\n",
    "        # define the forward computation on the latent z\n",
    "\n",
    "        # return the parameter for the output Bernoulli\n",
    "        # each is of size batch_size x 784\n",
    "        logits_img = self.fc1(z)\n",
    "\n",
    "        return logits_img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GBQ-eM3YAz4H",
    "tags": []
   },
   "source": [
    "# Training a VAE in Pyro"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JUcFqtTcfCra"
   },
   "source": [
    "### Setup training (single epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rVlgOvKlfCra",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(svi, data):\n",
    "\n",
    "    epoch_loss = svi.step(data)\n",
    "\n",
    "    # return epoch loss\n",
    "    normalizer_train = len(data)\n",
    "    total_epoch_loss_train = epoch_loss / normalizer_train\n",
    "    return total_epoch_loss_train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TJ2TYH7jfCrb"
   },
   "source": [
    "### Perform learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 626
    },
    "id": "4hDagpv8fCrb",
    "outputId": "2ca7ac38-d740-49dd-a6b4-e5e8bcbe37e2",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Dimensionality of latent space\n",
    "z_dim=2\n",
    "\n",
    "# Define the variational autoencoder\n",
    "vae = VAE(z_dim, encoder = Encoder(z_dim, hidden_dim=500), decoder = LinearDecoder(z_dim))\n",
    "\n",
    "# Run options\n",
    "LEARNING_RATE = 0.01\n",
    "\n",
    "# Number of learning epochs\n",
    "NUM_EPOCHS = 1000\n",
    "\n",
    "# clear param store\n",
    "pyro.clear_param_store()\n",
    "\n",
    "# setup the optimizer\n",
    "adam_args = {\"lr\": LEARNING_RATE}\n",
    "optimizer = Adam(adam_args)\n",
    "\n",
    "# setup the inference algorithm\n",
    "svi = SVI(vae.model, vae.guide, optimizer, loss=Trace_ELBO())\n",
    "train_elbo = []\n",
    "# training loop\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    total_epoch_loss_train = train(svi, data.data)\n",
    "    train_elbo.append(-total_epoch_loss_train)\n",
    "    if (epoch % 100) == 0:\n",
    "        print(\"[epoch %03d] average training loss: %.4f\" % (epoch, total_epoch_loss_train))    \n",
    "    if (epoch % 500) == 0:\n",
    "        vae.reconstruct_images()\n",
    "vae.reconstruct_images()\n",
    "plt.plot(range(len(train_elbo)), train_elbo)\n",
    "plt.xlabel(\"Number of iterations\")\n",
    "plt.ylabel(\"ELBO\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nixDOYYYfCrc"
   },
   "source": [
    "# Explore the data in the latent space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the representation of the data points in the latent space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_latent_representation(vae):\n",
    "    \"\"\"\n",
    "    Plot latent representations of data\n",
    "    \"\"\"\n",
    "    z_loc, z_scale = vae.encoder(data.data)\n",
    "    legends = [f\"Digit {i}\" for i in digits]\n",
    "    z_loc = z_loc.detach().numpy()\n",
    "\n",
    "    plt.figure(figsize=(5,5))\n",
    "    for idx, i in enumerate(digits):\n",
    "        plt.scatter(z_loc[data.targets.numpy()==i,0], z_loc[data.targets.numpy()==i,1], label=legends[idx])\n",
    "\n",
    "    plt.xlabel(r\"$Z_0$\")\n",
    "    plt.ylabel(r\"$Z_1$\")\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "plot_latent_representation(vae)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Partition the latent space and show the 'mean' image for each partition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_latent_space(vae: VAE, no_images: int = 15, img_size: int = 28):\n",
    "    \"\"\"\n",
    "    Plot the latent space of the VAE\n",
    "    \"\"\"\n",
    "    # Set up \"target\" z-values\n",
    "    # linearly spaced coordinates on the unit square were transformed\n",
    "    # through the inverse CDF (ppf) of the Gaussian to produce values\n",
    "    # of the latent variables z, since the prior of the latent space\n",
    "    # is Gaussian\n",
    "    z1 = norm.ppf(np.linspace(0.001, 0.999, no_images))\n",
    "    z2 = norm.ppf(np.linspace(0.999, 0.001, no_images)) # Flipped compared to z1 to make the figures comparable\n",
    "    \n",
    "    z_grid = np.dstack(np.meshgrid(z1, z2)).reshape(-1, 2)  # shape [no_images**2, z_dim]\n",
    "    x_pred_grid = vae.decoder.forward(\n",
    "        torch.tensor(z_grid, dtype=torch.float)\n",
    "    ).detach().numpy().reshape(no_images, no_images, img_size, img_size)\n",
    "    \n",
    "    images = dist.Bernoulli(logits=torch.tensor(x_pred_grid)).mean.detach().numpy()\n",
    "    \n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(np.block(list(map(list, images))), cmap='gray')\n",
    "    plt.gca().xaxis.set_major_locator(matplotlib.ticker.NullLocator())\n",
    "    plt.gca().yaxis.set_major_locator(matplotlib.ticker.NullLocator())\n",
    "    plt.box(False)\n",
    "    plt.show()\n",
    "    \n",
    "plot_latent_space(vae)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Show the posterior landscape of the latent space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_posterior(model):\n",
    "    z_samples,_ = model.encoder(data.data)\n",
    "    z_samples = z_samples.detach().numpy()\n",
    "    z_class = data.targets.detach().numpy()\n",
    "\n",
    "    for c in digits:\n",
    "        sns.kdeplot(data=pd.DataFrame(np.hstack((z_samples,z_class.reshape(-1,1))), columns=[\"dim1\", \"dim2\", \"Digit\"]), \n",
    "                    x=\"dim1\", y=\"dim2\", hue=\"Digit\", palette=\"crest\")  \n",
    "    plt.show()\n",
    "\n",
    "plot_posterior(vae)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cTDnjNrFBuhu"
   },
   "source": [
    "# Exercise\n",
    "\n",
    "1. Define as decoder network an MLP with one or more hidden layers and non-linear activation functions (e.g. Relu). \n",
    "    * You can also try to define a more expressive encoder network \n",
    "    * The number of hidden units should be parametrized.\n",
    "    * You can adopt/complete the classes *FancyDecoder* and *FancyEncoder* defined below, finding inspiration in the classes *Encoder* and *LinearDecoder* given above.\n",
    "2. Learn your newly defined VAE by adapting the code above. \n",
    "    * You may want to change the number of data points being used. A good starting point could be, e.g., 100 data points pr class while you experiment with the structure, but for getting more convincing visualization you may want to increase that to, say, 1000 data points pr. class\n",
    "3. Explore the latent space (both the latent representation of the data points as well as the structure of the latent space). What happens when we increase the capacity of the non-linear decoder (i.e we increase the number of hidden units). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Candidate solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class FancyDecoder(nn.Module):\n",
    "    def __init__(self, z_dim, hidden_dim1, hidden_dim2):\n",
    "        super(FancyDecoder, self).__init__()\n",
    "        \n",
    "        # Define the first layer\n",
    "        self.fc1 = nn.Linear(z_dim, hidden_dim1)\n",
    "  \n",
    "        # Define the first layer\n",
    "        self.fc2 = nn.Linear(hidden_dim1, hidden_dim2)\n",
    "        \n",
    "        # Define the last layer\n",
    "        self.fc3 = nn.Linear(hidden_dim2, 784)\n",
    "\n",
    "        # Define the non-linearities\n",
    "        self.non_linearity = nn.ReLU()\n",
    "\n",
    "    def forward(self, z):\n",
    "        # define the forward computation on the latent z\n",
    "\n",
    "        # first compute the hidden units\n",
    "        hidden = self.non_linearity(self.fc1(z))\n",
    "        hidden = self.non_linearity(self.fc2(hidden))\n",
    "        \n",
    "        # return the parameter for the output Bernoulli\n",
    "        # each is of size batch_size x 784\n",
    "        logits_img = self.fc3(hidden)\n",
    "\n",
    "        return logits_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class FancyEncoder(nn.Module):\n",
    "    def __init__(self, z_dim, hidden_dim1, hidden_dim2):\n",
    "        super(FancyEncoder, self).__init__()\n",
    "        # setup the three linear transformations used\n",
    "        self.fc1 = nn.Linear(784, hidden_dim1)\n",
    "        self.fc2 = nn.Linear(hidden_dim1, hidden_dim2)\n",
    "        self.fc31 = nn.Linear(hidden_dim2, z_dim)\n",
    "        self.fc32 = nn.Linear(hidden_dim2, z_dim)\n",
    "        # setup the non-linearities\n",
    "        self.non_linearity = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # define the forward computation on the image x\n",
    "        # first shape the mini-batch to have pixels in the rightmost dimension\n",
    "        x = x.reshape(-1, 784)\n",
    "        # then compute the hidden units\n",
    "        hidden = self.non_linearity(self.fc1(x))\n",
    "        hidden = self.non_linearity(self.fc2(hidden))\n",
    "        # then return a mean vector and a (positive) square root covariance\n",
    "        # each of size batch_size x z_dim\n",
    "        z_loc = self.fc31(hidden)\n",
    "        z_scale =nn.functional.softplus(self.fc32(hidden))\n",
    "        return z_loc, z_scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Uncomment the following line if you want to play around with different data sets\n",
    "#data, digits, number_of_images_from_each_class = get_data_sample(number_of_images_from_each_class = 1000)\n",
    "\n",
    "z_dim=2\n",
    "\n",
    "vae_non_linear = VAE(z_dim, encoder = FancyEncoder(z_dim, hidden_dim1 = 256, hidden_dim2 = 64), decoder = FancyDecoder(z_dim, hidden_dim1=64, hidden_dim2=256))\n",
    "\n",
    "# Run options\n",
    "LEARNING_RATE = 0.001\n",
    "\n",
    "# Number of learning epochs\n",
    "NUM_EPOCHS = 1000\n",
    "\n",
    "# clear param store\n",
    "pyro.clear_param_store()\n",
    "\n",
    "# setup the optimizer\n",
    "adam_args = {\"lr\": LEARNING_RATE}\n",
    "optimizer = Adam(adam_args)\n",
    "\n",
    "# setup the inference algorithm\n",
    "svi = SVI(vae_non_linear.model, vae_non_linear.guide, optimizer, loss=Trace_ELBO())\n",
    "train_elbo = []\n",
    "# training loop\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    total_epoch_loss_train = train(svi, data.data)\n",
    "    train_elbo.append(-total_epoch_loss_train)\n",
    "    if (epoch % 100) == 0:\n",
    "        print(\"[epoch %03d] average training loss: %.4f\" % (epoch, total_epoch_loss_train))\n",
    "    \n",
    "    if (epoch % 500) == 0:\n",
    "        vae_non_linear.reconstruct_images()\n",
    "vae_non_linear.reconstruct_images()\n",
    "\n",
    "plt.plot(range(len(train_elbo)), train_elbo)\n",
    "plt.xlabel(\"Number of iterations\")\n",
    "plt.ylabel(\"ELBO\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explore the latent space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualizing the data points in latent space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_latent_representation(vae)\n",
    "plot_latent_representation(vae_non_linear)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Representation of the latent space. Linearly spaced coordinates on the unit sphere is mapped though the inverse of the inverse of the commulative density function of the Gaussian to find values $\\boldsymbol z$ in latent space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Using the linear VAE\n",
    "plot_latent_space(vae)\n",
    "# Using the fancy VAE\n",
    "plot_latent_space(vae_non_linear)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_posterior(vae)\n",
    "plot_posterior(vae_non_linear)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Reconstruct images using the VAE\n",
    "vae.reconstruct_images()\n",
    "# Reconstruct images using the non-linear VAE\n",
    "vae_non_linear.reconstruct_images()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "name": "solutions_VAE.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}