notebook/notebook.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "VAE.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SH8CShRYPYhv"
      },
      "source": [
        "#Installations & Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LBLVaAuonBIw"
      },
      "source": [
        "!pip install pytorch-lightning"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lf1HdB31O5oq"
      },
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import pytorch_lightning as pl\n",
        "import numpy as np\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KMBJm1iWV2m9"
      },
      "source": [
        "#Parameter Initialization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qq6avijbmPYw"
      },
      "source": [
        "import os\n",
        "# Insert your own path here\n",
        "ds_path = os.path.join(\".\", \"drive\", \"My Drive\", \"Machine Learning\", \"Datasets\")\n",
        "\n",
        "variation = \"B-VAE\"\n",
        "\n",
        "configuration = {\n",
        "    \"dataset\": \"MNIST\",\n",
        "    \"path\": ds_path\n",
        "}\n",
        "\n",
        "architecture = {\n",
        "    \"conv_layers\": 3,\n",
        "    \"conv_channels\": [128, 256, 512],\n",
        "    \"conv_kernel_sizes\": [(7, 7), (7, 7), (5, 5)],\n",
        "    \"conv_strides\": [(1, 1), (1, 1), (1, 1)],\n",
        "    \"conv_paddings\": [(1, 1), (1, 1), (1, 1)],\n",
        "    \"z_dimension\": 2\n",
        "}\n",
        "\n",
        "hyperparameters = {\n",
        "    \"epochs\": 10,\n",
        "    \"batch_size\": 16,\n",
        "    \"learning_rate\": 3e-6,\n",
        "    \"beta\": 150\n",
        "}"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P_6ejOn_V-YN"
      },
      "source": [
        "def prepare_dataset(configuration):\n",
        "    \"\"\"\n",
        "    :param dict configuration: The configuration dictionary returned by parse_config_file\n",
        "\n",
        "    :return:        A dictionary containing information about the dataset used\n",
        "\n",
        "    Function used to set some values used by the model based on the dataset selected\n",
        "    \"\"\"\n",
        "    dataset_info = {}\n",
        "    if (configuration[\"dataset\"] == \"MNIST\"):\n",
        "        dataset_info[\"ds_method\"] = torchvision.datasets.MNIST\n",
        "        dataset_info[\"ds_shape\"] = (1, 28, 28)\n",
        "        dataset_info[\"ds_path\"] = configuration[\"path\"]\n",
        "    elif (configuration[\"dataset\"] == \"CIFAR10\"):\n",
        "        dataset_info[\"ds_method\"] = torchvision.datasets.CIFAR10\n",
        "        dataset_info[\"ds_shape\"] = (3, 32, 32)\n",
        "        dataset_info[\"ds_path\"] = configuration[\"path\"]\n",
        "    elif (configuration[\"dataset\"] == \"FashionMNIST\"):\n",
        "        dataset_info[\"ds_method\"] = torchvision.datasets.FashionMNIST\n",
        "        dataset_info[\"ds_shape\"] = (1, 28, 28)\n",
        "        dataset_info[\"ds_path\"] = configuration[\"path\"]\n",
        "    else:\n",
        "        print(\"Currently only MNIST & CIFAR10 datasets are supported\")\n",
        "        return None\n",
        "\n",
        "    return dataset_info"
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G-eW2P6roRBo"
      },
      "source": [
        "dataset_info = prepare_dataset(configuration)"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4wcGIiNyWGre"
      },
      "source": [
        "#Model Declaration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S8n44hIDWRcm"
      },
      "source": [
        "#####Utility Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RlfcNIqBWTW4"
      },
      "source": [
        "from torch import nn\n",
        "def compute_output_shape(current_shape, kernel_size, stride, padding):\n",
        "    \"\"\"\n",
        "    :param tuple current_shape:  The current shape of the data before a convolution is applied.\n",
        "    :param tuple kernel_size:    The kernel size of the current convolution operation.\n",
        "    :param tuple stride:         The stride of the current convolution operation.\n",
        "    :param tuple padding:        The padding of the current convolution operation.\n",
        "\n",
        "    :return:  The shape after a convolution operation with the above parameters is applied.\n",
        "    :rtype:   tuple\n",
        "\n",
        "            The formula used to compute the final shape is\n",
        "\n",
        "        component[i] = floor((N[i] - K[i] + 2 * P[i]) / S[i]) + 1\n",
        "\n",
        "        where, N = current shape of the data\n",
        "               K = kernel size\n",
        "               P = padding\n",
        "               S = stride\n",
        "    \"\"\"\n",
        "    # get the dimension of the data compute each component using the above formula\n",
        "    dimensions = len(current_shape)\n",
        "    return tuple((current_shape[i] - kernel_size[i] + 2 * padding[i]) // stride[i] + 1\n",
        "                 for i in range(dimensions))\n",
        "\n",
        "\n",
        "def compute_transpose_output_shape(current_shape, kernel_size, stride, padding):\n",
        "    \"\"\"\n",
        "    :param tuple current_shape:  The current shape of the data before a transpose convolution is\n",
        "                                   applied.\n",
        "    :param tuple kernel_size:    The kernel size of the current transpose convolution operation.\n",
        "    :param tuple stride:         The stride of the current transpose convolution operation.\n",
        "    :param tuple padding:        The padding of the current transpose convolution operation.\n",
        "\n",
        "    :return:  The shape after a transpose convolution operation with the above parameters is\n",
        "                applied.\n",
        "    :rtype:   tuple\n",
        "\n",
        "            The formula used to compute the final shape is\n",
        "\n",
        "        component[i] = (N[i] - 1) * S[i] - 2 * P[i] + (K[i] - 1) + 1\n",
        "\n",
        "        where, N = current shape of the data\n",
        "               K = kernel size\n",
        "               P = padding\n",
        "               S = stride\n",
        "    \"\"\"\n",
        "    # get the dimension of the data compute each component using the above formula\n",
        "    dimensions = len(current_shape)\n",
        "    return tuple((current_shape[i] - 1) * stride[i] - 2 * padding[i] + (kernel_size[i] - 1) + 1\n",
        "                 for i in range(dimensions))\n",
        "\n",
        "\n",
        "def compute_output_padding(current_shape, target_shape):\n",
        "    \"\"\"\n",
        "    :param tuple current_shape:  The shape of the data after a transpose convolution operation\n",
        "                                   takes place.\n",
        "    :param tuple target_shape:   The target shape that we would like our data to have after the\n",
        "                                   transpose convolution operation takes place.\n",
        "\n",
        "    :return:  The output padding needed so that the shape of the image after a transpose\n",
        "                convolution is applied, is the same as the target shape.\n",
        "    :rtype:   tuple\n",
        "    \"\"\"\n",
        "    # basically subtract each term to get the difference which will be the output padding\n",
        "    dimensions = len(current_shape)\n",
        "    return tuple(target_shape[i] - current_shape[i] for i in range(dimensions))\n",
        "\n",
        "\n",
        "def invalid_shape(current_shape):\n",
        "    \"\"\"\n",
        "    :param tuple current_shape:  The current shape of the data after a convolution is applied.\n",
        "\n",
        "    :return:  True if the shape is invalid, that is, a negative or 0 components exists. Else, it\n",
        "                returns False.\n",
        "    :rtype:   bool\n",
        "    \"\"\"\n",
        "    # check all components\n",
        "    for component in current_shape:\n",
        "        if component <= 0:\n",
        "            return True\n",
        "    # return False if they are ok\n",
        "    return False\n",
        "\n",
        "\n",
        "def create_encoder(architecture, input_shape):\n",
        "    \"\"\"\n",
        "    :param dict architecture:  A dictionary containing the hyperparameters that define the\n",
        "                                 architecture of the model.\n",
        "    :param tuple input_shape:  A tuple that corresponds to the shape of the input.\n",
        "\n",
        "    :return:  A PyTorch Sequential model that represents the encoder part of a VAE, along with the\n",
        "                final shape that a data point would have after the sequential is applied to it.\n",
        "    :rtype:   (torch.nn.Sequential, tuple)\n",
        "\n",
        "    This method builds the encoder part of a VAE and returns it. It is common for all types of VAE.\n",
        "    \"\"\"\n",
        "\n",
        "    # initialize useful variables\n",
        "    in_channels = input_shape[0]\n",
        "    current_shape = (input_shape[1], input_shape[2])\n",
        "\n",
        "    # initialize a list that will store the shape produced in each layer\n",
        "    shape_per_layer = [current_shape]\n",
        "\n",
        "    # build the encoder part\n",
        "    conv_sets = []\n",
        "\n",
        "    # iterate through the lists that define the architecture of the encoder\n",
        "    for layer in range(architecture[\"conv_layers\"]):\n",
        "\n",
        "        # get the variables from the dictionary for more verbose\n",
        "        out_channels = architecture[\"conv_channels\"][layer]\n",
        "        kernel_size = architecture[\"conv_kernel_sizes\"][layer]\n",
        "        stride = architecture[\"conv_strides\"][layer]\n",
        "        padding = architecture[\"conv_paddings\"][layer]\n",
        "\n",
        "        # add a set of Convolutional - Leaky ReLU - Batch Normalization sequential layers\n",
        "        conv_sets.append(\n",
        "            nn.Sequential(\n",
        "                nn.Conv2d(in_channels=in_channels,\n",
        "                          out_channels=out_channels,\n",
        "                          kernel_size=kernel_size,\n",
        "                          stride=stride,\n",
        "                          padding=padding),\n",
        "                nn.LeakyReLU(negative_slope=0.15),\n",
        "                nn.BatchNorm2d(out_channels))\n",
        "        )\n",
        "\n",
        "        # compute the new shape of the image\n",
        "        current_shape = compute_output_shape(current_shape=current_shape,\n",
        "                                             kernel_size=kernel_size,\n",
        "                                             stride=stride,\n",
        "                                             padding=padding)\n",
        "        shape_per_layer.append(current_shape)\n",
        "\n",
        "        # make sure that the shape is valid, and if not, raise an error\n",
        "        if invalid_shape(current_shape):\n",
        "            raise InvalidArchitectureError(shape=current_shape, layer=layer+1)\n",
        "\n",
        "        # the output channels of the current layer becomes the input channels of the next layer\n",
        "        in_channels = out_channels\n",
        "\n",
        "    # create a Sequential model and return it (* asterisk is used to unpack the list)\n",
        "    return nn.Sequential(*conv_sets), shape_per_layer\n",
        "\n",
        "\n",
        "def create_decoder(architecture, encoder_shapes):\n",
        "    \"\"\"\n",
        "    :param dict architecture:    A dictionary containing the hyperparameters that define the\n",
        "                                   architecture of the model.\n",
        "    :param list encoder_shapes:  A list that contains the shape of the data after it is applied to\n",
        "                                    every set of convolutional layers.\n",
        "\n",
        "    :return:  A PyTorch Sequential model that represents the decoder part of a VAE.\n",
        "    :rtype:   (torch.nn.Sequential)\n",
        "\n",
        "    This method builds the decoder part of a VAE and returns it. It is common for all types of VAE.\n",
        "    \"\"\"\n",
        "    # now start building the decoder part\n",
        "    conv_sets = []\n",
        "\n",
        "    # initialize useful variables\n",
        "    in_channels = architecture[\"conv_channels\"][-1]\n",
        "\n",
        "    # iterate through the lists that define the architecture of the decoder\n",
        "    for layer in range(architecture[\"conv_layers\"] - 1, -1, -1):\n",
        "\n",
        "        # get the variables from the dictionary for more verbose\n",
        "        out_channels = architecture[\"conv_channels\"][layer]\n",
        "        kernel_size = architecture[\"conv_kernel_sizes\"][layer]\n",
        "        stride = architecture[\"conv_strides\"][layer]\n",
        "        padding = architecture[\"conv_paddings\"][layer]\n",
        "\n",
        "        # compute the output shape after a transpose convolution in order to get the output padding\n",
        "        current_shape = encoder_shapes[layer + 1]\n",
        "        target_shape = encoder_shapes[layer]\n",
        "        output_shape = compute_transpose_output_shape(current_shape=current_shape,\n",
        "                                                      kernel_size=kernel_size,\n",
        "                                                      stride=stride,\n",
        "                                                      padding=padding)\n",
        "        output_padding = compute_output_padding(output_shape, target_shape)\n",
        "\n",
        "        # add a set of ConvolutionalTranspose - Leaky ReLU - Batch Normalization sequential layers\n",
        "        conv_sets.append(\n",
        "            nn.Sequential(\n",
        "                nn.ConvTranspose2d(in_channels=in_channels,\n",
        "                                   out_channels=out_channels,\n",
        "                                   kernel_size=kernel_size,\n",
        "                                   stride=stride,\n",
        "                                   padding=padding,\n",
        "                                   output_padding=output_padding),\n",
        "                nn.LeakyReLU(negative_slope=0.15),\n",
        "                nn.BatchNorm2d(out_channels))\n",
        "        )\n",
        "\n",
        "        # the output channels of the current layer becomes the input channels of the next layer\n",
        "        in_channels = out_channels\n",
        "\n",
        "    # create a Sequential model and return it (* asterisk is used to unpack the list)\n",
        "    return nn.Sequential(*conv_sets)\n",
        "\n",
        "\n",
        "def create_output_layer(architecture, input_shape):\n",
        "    \"\"\"\n",
        "    :param dict architecture:  A dictionary containing the hyperparameters that define the\n",
        "                                 architecture of the model.\n",
        "    :param tuple input_shape:  A tuple that corresponds to the shape of the input.\n",
        "\n",
        "    :return:  A PyTorch Sequential model that represents the output layer of a VAE.\n",
        "    :rtype:   torch.nn.Sequential\n",
        "\n",
        "    This method creates the output layer of a VAE, that is, the layer where the data from the\n",
        "    output of the decoder gets fed in order to be finally reconstructed.\n",
        "    \"\"\"\n",
        "    # define the variables of the architecture for more verbose\n",
        "    in_channels = architecture[\"conv_channels\"][0]\n",
        "    kernel_size = architecture[\"conv_kernel_sizes\"][0]\n",
        "    stride = architecture[\"conv_strides\"][0]\n",
        "    padding = architecture[\"conv_paddings\"][0]\n",
        "\n",
        "    return nn.Sequential(nn.ConvTranspose2d(in_channels=in_channels,\n",
        "                                            out_channels=in_channels,\n",
        "                                            kernel_size=kernel_size,\n",
        "                                            stride=stride,\n",
        "                                            padding=padding),\n",
        "                         nn.SELU(),\n",
        "                         nn.BatchNorm2d(in_channels),\n",
        "                         nn.Conv2d(in_channels=in_channels,\n",
        "                                   out_channels=input_shape[0],\n",
        "                                   kernel_size=kernel_size,\n",
        "                                   stride=stride,\n",
        "                                   padding=padding),\n",
        "                         nn.Sigmoid())"
      ],
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bqv8srojWLdA"
      },
      "source": [
        "###Variational Auto-Encoder"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OJePIeJfWKqF"
      },
      "source": [
        "import multiprocessing\n",
        "\n",
        "class VAE(pl.LightningModule):\n",
        "    \"\"\" Class that implements a Variational Autoencoder \"\"\"\n",
        "\n",
        "    def __init__(self, architecture, hyperparameters, dataset_info):\n",
        "        \"\"\"\n",
        "        :param dict architecture:      A dictionary containing the hyperparameters that define the\n",
        "                                         architecture of the model.\n",
        "        :param dict hyperparameters:   A tuple that corresponds to the shape of the input.\n",
        "        :param dict dataset_info:      The dimension of the latent vector z (bottleneck).\n",
        "\n",
        "        The constructor of the Variational Autoencoder.\n",
        "        \"\"\"\n",
        "\n",
        "        # call the constructor of the super class\n",
        "        super(VAE, self).__init__()\n",
        "\n",
        "        # initialize class variables regarding the architecture of the model\n",
        "        self.conv_layers = architecture[\"conv_layers\"]\n",
        "        self.conv_channels = architecture[\"conv_channels\"]\n",
        "        self.conv_kernel_sizes = architecture[\"conv_kernel_sizes\"]\n",
        "        self.conv_strides = architecture[\"conv_strides\"]\n",
        "        self.conv_paddings = architecture[\"conv_paddings\"]\n",
        "        self.z_dim = architecture[\"z_dimension\"]\n",
        "\n",
        "        # unpack the \"hyperparameters\" dictionary\n",
        "        self.batch_size = hyperparameters[\"batch_size\"]\n",
        "        self.learning_rate = hyperparameters[\"learning_rate\"]\n",
        "        self.scheduler_step_size = hyperparameters[\"epochs\"] // 2\n",
        "\n",
        "        # unpack the \"dataset_info\" dictionary\n",
        "        self.dataset_method = dataset_info[\"ds_method\"]\n",
        "        self.dataset_shape = dataset_info[\"ds_shape\"]\n",
        "        self.dataset_path = dataset_info[\"ds_path\"]\n",
        "\n",
        "        # build the encoder\n",
        "        self.encoder, self.encoder_shapes = create_encoder(architecture, self.dataset_shape)\n",
        "\n",
        "        # compute the length of the output of the decoder once it has been flattened\n",
        "        in_features = self.conv_channels[-1] * np.prod(self.encoder_shapes[-1][:])\n",
        "        # now define the mean and standard deviation layers\n",
        "        self.mean_layer = nn.Linear(in_features=in_features, out_features=self.z_dim)\n",
        "        self.std_layer = nn.Linear(in_features=in_features, out_features=self.z_dim)\n",
        "\n",
        "        # use a linear layer for the input of the decoder\n",
        "        in_channels = self.conv_channels[-1]\n",
        "        self.decoder_input = nn.Linear(in_features=self.z_dim, out_features=in_features)\n",
        "\n",
        "        # build the decoder\n",
        "        self.decoder = create_decoder(architecture, self.encoder_shapes)\n",
        "\n",
        "        # build the output layer\n",
        "        self.output_layer = create_output_layer(architecture, self.dataset_shape)\n",
        "\n",
        "    def _encode(self, X):\n",
        "        \"\"\"\n",
        "        :param Tensor X:  Input to encode into mean and standard deviation. (N, C, H, W)\n",
        "\n",
        "        :return:  The mean and std tensors that the encoder produces for input X. (N, z_dim)\n",
        "        :rtype:   (Tensor, Tensor)\n",
        "\n",
        "        This method applies forward propagation to the self.encoder in order to get the mean and\n",
        "        standard deviation of the latent vector z.\n",
        "        \"\"\"\n",
        "        # run the input through the encoder part of the Network\n",
        "        encoded_input = self.encoder(X)\n",
        "\n",
        "        # flatten so that it can be fed to the mean and standard deviation layers\n",
        "        encoded_input = torch.flatten(encoded_input, start_dim=1)\n",
        "\n",
        "        # compute the mean and standard deviation\n",
        "        mean = self.mean_layer(encoded_input)\n",
        "        std = self.std_layer(encoded_input)\n",
        "\n",
        "        return mean, std\n",
        "\n",
        "    def _compute_latent_vector(self, mean, std):\n",
        "        \"\"\"\n",
        "        :param Tensor mean:  The mean of the latent vector z following a Gaussian distribution.\n",
        "                               (N, z_dim)\n",
        "        :param Tensor std:   The standard deviation of the latent vector z following a Gaussian\n",
        "                               distribution. (N, z_dim)\n",
        "\n",
        "        :return:  The Linear combination of the mean and standard deviation, where the latter\n",
        "                    factor is multiplied with a random variable epsilon ~ N(0, 1). Basically\n",
        "                    the latent vector z. (N, z_dim)\n",
        "        :rtype:   Tensor\n",
        "\n",
        "        This method computes the latent vector z by applying the reparameterization trick to the\n",
        "        output of the mean and standard deviation layers, in order to be able to later compute the\n",
        "        gradient. The stochasticiy here is introduced by the factor epsilon, which is an independent\n",
        "        node. Thus, we do not have to compute its gradient during backpropagation.\n",
        "        \"\"\"\n",
        "\n",
        "        # compute the stochastic node epsilon\n",
        "        epsilon = torch.randn_like(std)\n",
        "\n",
        "        # compute the linear combination of the above attributes and return\n",
        "        return mean + epsilon * (1.0 / 2) * std\n",
        "\n",
        "    def _decode(self, z):\n",
        "        \"\"\"\n",
        "        :param Tensor z:  Latent vector computed using the mean and variance layers (with the\n",
        "                            reparameterization trick). (N, z_dim)\n",
        "\n",
        "        :return:  The output of the decoder part of the network. (N, C, H, W)\n",
        "        :rtype:   Tensor\n",
        "\n",
        "        This method performs forward propagation of the latent vector through the decoder of the\n",
        "        VAE to get the final output of the network.\n",
        "        \"\"\"\n",
        "        # run the latent vector through the \"input decoder\" layer\n",
        "        decoder_input = self.decoder_input(z)\n",
        "\n",
        "        # convert back the shape that will be fed to the decoder\n",
        "        height = self.encoder_shapes[-1][0]\n",
        "        width = self.encoder_shapes[-1][1]\n",
        "        decoder_input = decoder_input.view(-1, self.conv_channels[-1], height, width)\n",
        "\n",
        "        # run through the decoder\n",
        "        decoder_output = self.decoder(decoder_input)\n",
        "\n",
        "        # run through the output layer and return\n",
        "        network_output = self.output_layer(decoder_output)\n",
        "        return network_output\n",
        "\n",
        "    def forward(self, X):\n",
        "        \"\"\"\n",
        "        :param Tensor X:  The input to run through the VAE. (N, C, H, W)\n",
        "\n",
        "        :return: The output of the Network, along with the mean, standard deviation layers.\n",
        "                   (N, C, H, W), (N, z_dim), (N, z_dim)\n",
        "        :rtype:  (Tensor, Tensor, Tensor)\n",
        "\n",
        "        This method performs Forward Propagation through all the layers of the VAE and returns\n",
        "        the reconstructed input.\n",
        "        \"\"\"\n",
        "        # encode the input to get mean and standard deviation\n",
        "        mean, std = self._encode(X)\n",
        "\n",
        "        # get the latent vector z by using the reparameterization trick\n",
        "        z = self._compute_latent_vector(mean, std)\n",
        "\n",
        "        # compute the output by propagating the latent vector through the decoder and return\n",
        "        decoded_output = self._decode(z)\n",
        "        return decoded_output, mean, std\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        \"\"\"\n",
        "        :return:  The optimizer that will be used during backpropagation, along with a scheduler.\n",
        "        :rtype:   (torch.optim.Optimizer, torch.optim.lr_scheduler)\n",
        "        \"\"\"\n",
        "        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)\n",
        "        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=self.scheduler_step_size,\n",
        "                                                    gamma=0.1)\n",
        "\n",
        "        return [optimizer], [scheduler]\n",
        "\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        \"\"\"\n",
        "        :param Tensor batch:   The current batch of the training set.\n",
        "        :param int batch_idx:  The batch index of the current batch.\n",
        "\n",
        "        :return:  A dictionary of the losses computed on the current prediction.\n",
        "        :rtype:   (dict)\n",
        "        \"\"\"\n",
        "        # unpack the current batch\n",
        "        X, y = batch\n",
        "\n",
        "        # pass it through the model\n",
        "        X_hat, mean, std = self(X)\n",
        "        # calculate the losses\n",
        "        losses = VAE.criterion(X, X_hat, mean, std)\n",
        "\n",
        "        return losses\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        \"\"\"\n",
        "        :return:  A DataLoader object of the training set.\n",
        "        :rtype:   torch.utils.data.DataLoader\n",
        "        \"\"\"\n",
        "        # download the training set using torchvision if it hasn't already been downloaded\n",
        "        train_set = self.dataset_method(root=self.dataset_path, train=True, download=True,\n",
        "                                        transform=torchvision.transforms.ToTensor())\n",
        "        # initialize a pytorch DataLoader to feed training batches into the model\n",
        "        self.train_loader = DataLoader(dataset=train_set, batch_size=self.batch_size, shuffle=True,\n",
        "                                       num_workers=multiprocessing.cpu_count()//2)\n",
        "        return self.train_loader\n",
        "\n",
        "    def test_step(self, batch, batch_idx):\n",
        "        \"\"\"\n",
        "        :param Tensor batch:   The current batch of the test set.\n",
        "        :param int batch_idx:  The batch index of the current batch.\n",
        "\n",
        "        :return:  A tuple consisting of a dictionary of the losses computed on the current\n",
        "                    prediction and the MSE Loss compared to the original picture.\n",
        "        :rtype:  (dict, Tensor)\n",
        "        \"\"\"\n",
        "        # unpack the current batch\n",
        "        X, y = batch\n",
        "\n",
        "        # pass it through the model\n",
        "        X_hat, mean, std = self(X)\n",
        "        # calculate the losses\n",
        "        losses = VAE.criterion(X, X_hat, mean, std)\n",
        "\n",
        "        # also calculate the MSE loss\n",
        "        mse_loss_func = torch.nn.MSELoss()\n",
        "        mse_loss = mse_loss_func(X, X_hat)\n",
        "\n",
        "        self.log('mse_loss', mse_loss.item())\n",
        "        self.log('losses', losses)\n",
        "        return losses, mse_loss\n",
        "\n",
        "    def test_dataloader(self):\n",
        "        \"\"\"\n",
        "        :return:  A DataLoader object of the test set.\n",
        "        :rtype:   torch.utils.data.DataLoader\n",
        "        \"\"\"\n",
        "        # download the test set using torchvision if it is not already downloaded\n",
        "        test_set = self.dataset_method(root=self.dataset_path, train=False, download=True,\n",
        "                                       transform=torchvision.transforms.ToTensor())\n",
        "        # initialize a pytorch DataLoader to feed test batches into the model\n",
        "        self.test_loader = DataLoader(dataset=test_set, batch_size=self.batch_size, shuffle=True,\n",
        "                                      num_workers=multiprocessing.cpu_count() // 2)\n",
        "        return self.test_loader\n",
        "\n",
        "    def sample(self, n):\n",
        "        \"\"\"\n",
        "        :param int n: The amount of images per row to sample\n",
        "\n",
        "        :return:  None\n",
        "\n",
        "        This method plots n^2 sampled images next to each other\n",
        "        \"\"\"\n",
        "        # create a new latent vector consisting of random values\n",
        "        z = torch.randn(n*n, self.z_dim)\n",
        "\n",
        "        # pass the vector through the decoder\n",
        "        samples = self._decode(z)\n",
        "\n",
        "        # set the correct colourmap that corresponds to the image dimension\n",
        "        cmap = None\n",
        "        if (self.dataset_shape[0] == 3):\n",
        "            cmap = 'viridis'\n",
        "        elif (self.dataset_shape[0] == 1):\n",
        "            cmap = 'gray'\n",
        "\n",
        "        plot_multiple(samples.detach().numpy(), n, self.dataset_shape, cmap)\n",
        "\n",
        "    def reconstruct(self, n):\n",
        "        \"\"\"\n",
        "        :param int n:  The number of images to plot in each row of whole plot.\n",
        "\n",
        "        :return:  None\n",
        "\n",
        "        This method plots n^2 reconstructed (from the test set) images next to each other.\n",
        "        \"\"\"\n",
        "\n",
        "        # get as many batches from the test set to fill the final plot\n",
        "        tensors = []\n",
        "        img_count = 0\n",
        "        while n * n > img_count:\n",
        "            batch, y = next(iter(self.test_loader))\n",
        "            img_count += len(batch)\n",
        "            tensors.append(batch)\n",
        "\n",
        "        # concatenate them\n",
        "        X = torch.cat(tensors, dim=0)\n",
        "\n",
        "        # pass them through the model\n",
        "        X_hat, mean, std = self(X)\n",
        "        min_imgs = min(n, len(X))\n",
        "\n",
        "        # set the correct colourmap that corresponds to the image dimension\n",
        "        cmap = None\n",
        "        if (self.dataset_shape[0] == 3):\n",
        "            cmap = 'viridis'\n",
        "        elif (self.dataset_shape[0] == 1):\n",
        "            cmap = 'gray'\n",
        "\n",
        "        # plot the images and their reconstructions\n",
        "        plot_multiple(X_hat.detach().numpy(), min_imgs, self.dataset_shape, cmap)\n",
        "\n",
        "    @staticmethod\n",
        "    def _data_fidelity_loss(X, X_hat, eps=1e-10):\n",
        "        \"\"\"\n",
        "        :param Tensor X:      The original input data that was passed to the VAE.\n",
        "                                (N, C, H, W)\n",
        "        :param Tensor X_hat:  The reconstructed data, the output of the VAE.\n",
        "                                (N, C, H, W)\n",
        "        :param Double eps:    A small positive double used to ensure we don't get log of 0.\n",
        "\n",
        "        :return:  A tensor containing the Data Fidelity term of the loss function,\n",
        "                    which is given by the formula below.\n",
        "        :rtype:   Tensor\n",
        "\n",
        "        E_{z ~ q_{phi}(z | x)}[log(p_{theta}(x|z))] = sum(x * log(x_hat) + (1 - x) * log(1 - x_hat))\n",
        "\n",
        "            which is basically a Cross Entropy Loss.\n",
        "\n",
        "        This method computes the Data Fidelity term of the loss function. A small positive double\n",
        "        epsilon is added inside the logarithm to make sure that we don't get log(0).\n",
        "        \"\"\"\n",
        "        # compute the data fidelity for every training example\n",
        "        data_fidelity = torch.sum(X * torch.log(eps + X_hat) + (1 - X) * torch.log(eps + 1 - X_hat),\n",
        "                                  axis=[1, 2, 3])\n",
        "        return data_fidelity\n",
        "\n",
        "    @staticmethod\n",
        "    def _kl_divergence_loss(mean, std):\n",
        "        \"\"\"\n",
        "        :param Tensor mean:  The output of the mean layer, computed with the output of the\n",
        "                               encoder. (N, z_dim)\n",
        "        :param Tensor std:   The output of the standard deviation layer, computed with the output\n",
        "                               of the encoder. (N, z_dim)\n",
        "\n",
        "        :return:  A tensor consisting of the KL-Divergence term of the loss function, which is\n",
        "                    given by the formula below.\n",
        "        :rtype:   Tensor\n",
        "\n",
        "        D_{KL}[q_{phi}(z | x) || p_{theta}(x)] = (1/2) * sum(std + mean^2 - 1 - log(std))\n",
        "\n",
        "            In the above equation we substitute std with e^{std} to improve numerical stability.\n",
        "\n",
        "        This method computes the KL-Divergence term of the loss function. It substitutes the\n",
        "        value of the standard deviation layer with exp(standard deviation) in order to ensure\n",
        "        numerical stability.\n",
        "        \"\"\"\n",
        "        # compute the kl divergence for each training example and return it\n",
        "        kl_divergence = (1 / 2) * torch.sum(torch.exp(std) + torch.square(mean) - 1 - std, axis=1)\n",
        "        return kl_divergence\n",
        "\n",
        "    @staticmethod\n",
        "    def criterion(X, X_hat, mean, std):\n",
        "        \"\"\"\n",
        "        :param Tensor X:      The original input data that was passed to the VAE.\n",
        "                                (N, C, H, W)\n",
        "        :param Tensor X_hat:  The reconstructed data, the output of the VAE.\n",
        "                                (N, C, H, W)\n",
        "        :param Tensor mean:   The output of the mean layer, computed with the output of the\n",
        "                                encoder. (N, z_dim)\n",
        "        :param Tensor std:    The output of the standard deviation layer, computed with the output\n",
        "                                of the encoder. (N, z_dim)\n",
        "\n",
        "        :return: A dictionary containing the values of the losses computed.\n",
        "        :rtype:  dict\n",
        "\n",
        "        This method computes the loss of the VAE using the formula:\n",
        "\n",
        "            L(x, x_hat) = - E_{z ~ q_{phi}(z | x)}[log(p_{theta}(x|z))]\n",
        "                          + D_{KL}[q_{phi}(z | x) || p_{theta}(x)]\n",
        "\n",
        "        Intuitively, the expectation term is the Data Fidelity term, and the second term is a\n",
        "        regularizer that makes sure the distribution of the encoder and the decoder stay close.\n",
        "        \"\"\"\n",
        "        # get the 2 losses\n",
        "        data_fidelity_loss = VAE._data_fidelity_loss(X, X_hat)\n",
        "        kl_divergence_loss = VAE._kl_divergence_loss(mean, std)\n",
        "\n",
        "        # add them to compute the loss for each training example in the mini batch\n",
        "        loss = -data_fidelity_loss + kl_divergence_loss\n",
        "\n",
        "        # place them all inside a dictionary and return it\n",
        "        losses = {\"data_fidelity\": torch.mean(data_fidelity_loss),\n",
        "                  \"kl-divergence\": torch.mean(kl_divergence_loss),\n",
        "                  \"loss\": torch.mean(loss)}\n",
        "        return losses"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M3YbhtuMXPTG"
      },
      "source": [
        "###B-VAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mzH1LDJ_XRfU"
      },
      "source": [
        "class betaVAE(VAE):\n",
        "    \"\"\"\n",
        "    Class that implements a Disentangled Variational Autoencoder (Beta VAE).\n",
        "    This class inherits from the VAE class.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, architecture, hyperparameters, dataset_info):\n",
        "        \"\"\"\n",
        "        :param dict architecture:     A dictionary containing the hyperparameters that define the\n",
        "                                        architecture of the model.\n",
        "        :param dict hyperparameters:  A tuple that corresponds to the shape of the input.\n",
        "        :param dict dataset_info:     The dimension of the latent vector z (bottleneck).\n",
        "\n",
        "        The constructor of the Disentangled Variational Autoencoder.\n",
        "        \"\"\"\n",
        "\n",
        "        # invoke the constructor of the VAE class, as the architecture is the same\n",
        "        super(betaVAE, self).__init__(architecture, hyperparameters, dataset_info)\n",
        "\n",
        "        # store the value of beta in the class as it exists only in this VAE variation\n",
        "        self.beta = hyperparameters[\"beta\"]\n",
        "\n",
        "    @staticmethod\n",
        "    def criterion(X, X_hat, mean, std):\n",
        "        \"\"\"\n",
        "        :param Tensor X:      The original input data that was passed to the B-VAE.\n",
        "                                (N, input_shape[1], H, W)\n",
        "        :param Tensor X_hat:  The reconstructed data, the output of the B-VAE.\n",
        "                                (N, input_shape[1], H, W)\n",
        "        :param Tensor mean:   The output of the mean layer, computed with the output of the\n",
        "                                encoder. (N, z_dim)\n",
        "        :param Tensor std:    The output of the standard deviation layer, computed with the output\n",
        "                                of the encoder. (N, z_dim)\n",
        "\n",
        "        :return:  A dictionary containing the values of the losses computed.\n",
        "        :rtype:   dict\n",
        "\n",
        "        This method computes the loss of the B-VAE using the formula:\n",
        "\n",
        "            L(x, x_hat) = - E_{z ~ q_{phi}(z | x)}[log(p_{theta}(x|z))]\n",
        "                          + beta * D_{KL}[q_{phi}(z | x) || p_{theta}(x)]\n",
        "\n",
        "        Intuitively, the expectation term is the Data Fidelity term, and the second term is a\n",
        "        regularizer that makes sure the distribution of the encoder and the decoder stay close.\n",
        "        \"\"\"\n",
        "        # get the 2 losses\n",
        "        data_fidelity_loss = VAE._data_fidelity_loss(X, X_hat)\n",
        "        kl_divergence_loss = VAE._kl_divergence_loss(mean, std)\n",
        "\n",
        "        # add them to compute the loss for each training example in the mini batch\n",
        "        loss = -data_fidelity_loss + self.beta * kl_divergence_loss\n",
        "\n",
        "        # place them all inside a dictionary and return it\n",
        "        losses = {\"data_fidelity\": torch.mean(data_fidelity_loss),\n",
        "                  \"kl-divergence\": torch.mean(kl_divergence_loss),\n",
        "                  \"beta_kl-divergence\": self.beta * torch.mean(kl_divergence_loss),\n",
        "                  \"loss\": torch.mean(loss)}\n",
        "        return losses"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H--U2gfVXdbc"
      },
      "source": [
        "#Model Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4j38U8JZXxTY"
      },
      "source": [
        "###Trainer Initialization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mUfz6Xu2iJzF"
      },
      "source": [
        "from pytorch_lightning import Trainer\n",
        "\n",
        "model = None\n",
        "if (variation == \"VAE\"):\n",
        "  model = VAE(architecture, hyperparameters, dataset_info)\n",
        "elif (variation == \"B-VAE\"):\n",
        "  model = betaVAE(architecture, hyperparameters, dataset_info)\n",
        "\n",
        "trainer = Trainer(max_epochs = hyperparameters[\"epochs\"], gpus=1, progress_bar_refresh_rate=20)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mVXcs2NdX59G"
      },
      "source": [
        "###Model fitting (training)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zQXiOhf6X9f-"
      },
      "source": [
        "trainer.fit(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2IVk_nbjYAxB"
      },
      "source": [
        "#Model testing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "te7gGiJTYGwH"
      },
      "source": [
        "###Testing on the test set"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "51AenrmDid1i"
      },
      "source": [
        "result = trainer.test(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gVRCAAxGYI39"
      },
      "source": [
        "###Utility function to plot two images against each other"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fxdiiG8UYFpH"
      },
      "source": [
        "def plot_multiple(images, n, dim, cmap):\n",
        "    \"\"\"\n",
        "    :param arr images:          An array of images stored as a numpy array\n",
        "    :param int n:               The width and height of the plot in terms of images\n",
        "    :param tuple dim:           The dimension of the images\n",
        "    :param str cmap:            The colourmap to be used by pyplot\n",
        "\n",
        "    :return:                    Nothing\n",
        "\n",
        "    Function used to plot multiple images in one single plot\n",
        "    \"\"\"\n",
        "    # unpack the image dimensions\n",
        "    z_dim, x_dim, y_dim = dim\n",
        "\n",
        "    # if image is grayscale\n",
        "    if (z_dim == 1):\n",
        "        # initialize some limits on x&y\n",
        "        x_limit = np.linspace(-2, 2, n)\n",
        "        y_limit = np.linspace(-2, 2, n)\n",
        "\n",
        "        # initialize the final combined image\n",
        "        empty = np.empty((x_dim*n, y_dim*n))\n",
        "\n",
        "        current = 0\n",
        "        for i, zi in enumerate(x_limit):\n",
        "            for j, pi in enumerate(y_limit):\n",
        "                # each image insert it into a subsection of the final image\n",
        "                empty[(n-i-1)*x_dim:(n-i)*x_dim, j*y_dim:(j+1)*y_dim] = images[current][0]\n",
        "                current+=1\n",
        "\n",
        "        plt.figure(figsize=(8, 10))\n",
        "\n",
        "        x,y = np.meshgrid(x_limit, y_limit)\n",
        "        plt.imshow(empty, origin=\"upper\", cmap=cmap)\n",
        "        plt.grid(False)\n",
        "        plt.show()\n",
        "\n",
        "    # if the image is rgb\n",
        "    elif (z_dim == 3):\n",
        "        # initialize some limits on x&y\n",
        "        x_limit = np.linspace(-2, 2, n)\n",
        "        y_limit = np.linspace(-2, 2, n)\n",
        "\n",
        "        # initialize the final combined image (now with one more dim)\n",
        "        empty = np.empty((x_dim*n, y_dim*n, 3))\n",
        "\n",
        "        current = 0\n",
        "        for i, zi in enumerate(x_limit):\n",
        "            for j, pi in enumerate(y_limit):\n",
        "                # flatten the image\n",
        "                curr_img = images[current].ravel()\n",
        "                # reshape it into the correct shape for pyplot\n",
        "                curr_img = np.reshape(curr_img, (x_dim, y_dim, z_dim), order='F')\n",
        "                # rotate it by 270 degrees\n",
        "                curr_img = np.rot90(curr_img, 3)\n",
        "\n",
        "                # insert it into a subsection of the final image\n",
        "                empty[(n-i-1)*x_dim:(n-i)*x_dim, j*y_dim:(j+1)*y_dim] = curr_img\n",
        "                current+=1\n",
        "\n",
        "        plt.figure(figsize=(8, 10))\n",
        "\n",
        "        x,y = np.meshgrid(x_limit, y_limit)\n",
        "        plt.imshow(empty, origin=\"upper\", cmap=cmap)\n",
        "        plt.grid(False)\n",
        "        plt.show()"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KDqPIPc4YQmn"
      },
      "source": [
        "###Small sample of the model to evaluate the results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xL6O3F-fihFd"
      },
      "source": [
        "model = model.cpu()\n",
        "model.sample(12)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xrNCbhvvJjmu"
      },
      "source": [
        "model.reconstruct(12)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}