08_autoencoders.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook was inspired by neural network & machine learning labs led by [GMUM](https://gmum.net/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See also [Chapter 14](https://www.deeplearningbook.org/contents/autoencoders.html) of the Deep Learning book, Lilian Weng's [From Autoencoder to Beta-VAE](https://lilianweng.github.io/lil-log/2018/08/12/from-autoencoder-to-beta-vae.html), and the [DCGAN Tutorial](https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Utils and imports (run and hide)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import Subset\n",
    "from torchvision.datasets import MNIST, FashionMNIST\n",
    "from torchvision.transforms import ToTensor, Lambda, Compose\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "def plot_dataset(train_data, model):\n",
    "    view_data = train_data.data[:5].view(-1, 28*28) / 255.\n",
    "    _, decoded_data = model.forward(train_data.data[:5].view(-1, 784).float().to(device) / 255.)\n",
    "    decoded_data = decoded_data.cpu().detach().numpy()\n",
    "\n",
    "    n_rows = 2 if decoded_data is not None else 1\n",
    "    n_cols = len(view_data)\n",
    "    plt.suptitle(\"Reconstruction\")\n",
    "    fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 2, n_rows * 2))\n",
    "    \n",
    "    if decoded_data is not None:\n",
    "        for i in range(n_cols):\n",
    "            axes[0][i].imshow(np.reshape(view_data.data.numpy()[i], (28, 28)), cmap='gray')\n",
    "            axes[0][i].set_xticks(())\n",
    "            axes[0][i].set_yticks(())\n",
    "        \n",
    "        for i in range(n_cols):\n",
    "            axes[1][i].clear()\n",
    "            axes[1][i].imshow(np.reshape(decoded_data[i], (28, 28)), cmap='gray')\n",
    "            axes[1][i].set_xticks(())\n",
    "            axes[1][i].set_yticks(())\n",
    "    \n",
    "    else:\n",
    "        for i in range(n_cols):\n",
    "            axes[i].imshow(np.reshape(view_data.data.numpy()[i], (28, 28)), cmap='gray')\n",
    "            axes[i].set_xticks(())\n",
    "            axes[i].set_yticks(())\n",
    "    \n",
    "    plt.show()\n",
    "    \n",
    "def plot_pca(data, model):\n",
    "    labels = data.classes\n",
    "    plt.suptitle(\"Reduction of latent space\")\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    pca = PCA(2)\n",
    "\n",
    "    z = model.encode(train_data.data.view(-1, 784).float().to(device))\n",
    "    reduced_z = pca.fit_transform(z.detach().cpu().numpy())\n",
    "    \n",
    "    for class_idx in range(10):\n",
    "        indices = (data.targets == class_idx)\n",
    "        plt.scatter(\n",
    "            reduced_z[indices, 0], reduced_z[indices, 1],\n",
    "            s=2., label=labels[class_idx])\n",
    "\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "    \n",
    "torch.manual_seed(1337) \n",
    "batch_size = 128 \n",
    "transforms = Compose([ToTensor(), Lambda(lambda x: x.flatten())])\n",
    "\n",
    "train_data = MNIST(root='.', \n",
    "                   train=True, \n",
    "                   transform=transforms,    \n",
    "                   download=True)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discriminative vs. generative models\n",
    "\n",
    "There are (generally) two approaches to statistical classification:\n",
    "- *generative*, where we model the joint probability distribution $P(X,Y)$,\n",
    "- *discriminative*, where we model the conditional probability of the target $Y$ given an observation $x$, $P(Y\\vert X=x)$ (classifiers computed without using a probability model are also loosely referred to as \"discriminative\").\n",
    "\n",
    "One can also think of generative models as learning the distribution of individual classes and discriminative models as learning (hard or soft) boundaries between classes.\n",
    "\n",
    "Generative models allow you to generate data similar to training data, whereas discriminative models might be easier to learn.\n",
    "\n",
    "See the [Wikipedia article](https://en.wikipedia.org/wiki/Generative_model) on generative models, the [CrossValidated question](https://stats.stackexchange.com/questions/12421/generative-vs-discriminative), and the classic ML paper [On Discriminative vs. Generative Classifiers: A comparison of logistic regression and naive Bayes](https://ai.stanford.edu/~ang/papers/nips01-discriminativegenerative.pdf) for more disambiguation.\n",
    "\n",
    "Today we will work with the autoencoder model, first showing how to use it for semi-supervised learning, and later building a generative model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vanilla autoencoder\n",
    "\n",
    "An *autoencoder* is a neural network that is trained to copy it's input to its output. The network may be viewed as consisting of two parts: an *encoder* $g_\\phi$, which takes in an input $\\mathbf{x}$ and produces a *code* (also: *hidden representation*; *latent vector*) $\\mathbf{z}=g_\\phi(\\mathbf{x})$ , and the *decoder* $f_\\theta$, which produces a reconstruction $\\mathbf{x'}=f_\\theta(\\mathbf{z})$.\n",
    "\n",
    "![auto-encoder](figures/ae.png)\n",
    "<center>Source: <a href=\"https://lilianweng.github.io/lil-log/2018/08/12/from-autoencoder-to-beta-vae.html\">From Autoencoder to Beta-VAE</a>.</center>\n",
    "\n",
    "If an autoencoder succeeds in simply learning to set $f_\\theta(g_\\phi(\\mathbf{x}))=\\mathbf{x}$ everywhere, then it is not especially useful. Instead, autoencoders are designed to be unable to learn to copy perfectly. Because the model is forced to prioritize which aspects of the input should be copied, it often learns useful properties of the data.\n",
    "\n",
    "The loss function for the vanilla autoencoder is the MSE between the input and output:\n",
    "$$L_{AE} =\\frac{1}{n}\\sum_i \\lVert\\mathbf{x}_i-f_\\theta(g_\\phi(\\mathbf{x}_i))\\rVert_2^2.$$\n",
    "The encoder and the decoder can be arbitrary neural networks, but usually the decoder is comprised of the same transformations as the encoder in reverse order.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1 (0.25p)\n",
    "Implement the encoder and the decoder for a vanilla autoencoder. \n",
    "\n",
    "The dimensions in the encoder are supposed to have the following number of neurons: `(784, 128, 128, 64, latent_dim)`. Analogously, for the decoder: `(latent_dim, 64, 128, 128, 784)`. (The input and output dimensionality corresponds to the number of pixels in MNIST.) `latent_dim` is supposed to be a parameter of the constructor.\n",
    "\n",
    "Use no activation function after the encoder, a sigmoid after the decoder, and ReLU after the hidden layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AutoEncoder(torch.nn.Module):\n",
    "    def __init__(self, latent_dim):\n",
    "        \n",
    "        super(AutoEncoder, self).__init__()\n",
    "        \n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        self.encoder = ???\n",
    "        self.decoder = ???\n",
    "    \n",
    "    def encode(self, x):\n",
    "        return ???\n",
    "    \n",
    "    def decode(self, encoded):\n",
    "        return ???\n",
    "\n",
    "    def forward(self, x):\n",
    "        encoded = ???\n",
    "        decoded = ???\n",
    "        return encoded, decoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 25\n",
    "lr = 5e-3        \n",
    "\n",
    "autoencoder = AutoEncoder(latent_dim=10).to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(autoencoder.parameters(), lr=lr)\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    epoch_losses = []  \n",
    "    for step, (x, _) in enumerate(train_loader):\n",
    "        \n",
    "        x = x.to(device)\n",
    "        \n",
    "        optimizer.zero_grad()        \n",
    "\n",
    "        _, decoded = autoencoder(x) \n",
    "        loss = criterion(decoded, x)   \n",
    "        loss.backward()          \n",
    "        optimizer.step()            \n",
    "        \n",
    "        epoch_losses.append(loss.item())\n",
    "\n",
    "    print(f'Epoch: {epoch+1}  |  train loss: {np.mean(epoch_losses):.4f}')\n",
    "\n",
    "    if epoch % 10 == 0:\n",
    "        plot_dataset(train_data, autoencoder)\n",
    "        plot_pca(train_data, autoencoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Semi-supervised learning\n",
    "\n",
    "In practice building a fully-labeled dataset can be very costly. If we want to train an image classifier, then gathering a large amount of data isn't a problem (we can scrape them from the internet, for example). Labelling them, however, is, and would require human resources. In some cases, labelling can be even more expensive -- in the segmentation task, where we want to assign a class to each pixel in the image, assigning the labels for one picture can take many hours.\n",
    "\n",
    "Thus, we'd like to have methods which are able to utilize data for which we don't have labels. In the following task we'll build a simple semi-supervised model using an autoencoder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2 (0.5p)\n",
    "\n",
    "Assume that for the 60k examples from MNIST only 100 have the label. The 100 labeled examples are in the variable `labeled_data`.\n",
    "\n",
    "1. Implement a classifier and train it on the 100 labeled examples. Report the accuracy on the test set. (The net should be relatively simple: max. 4 layers, max. 128 neurons in a layer).\n",
    "2. Implement a classifier and train it on the 100 labeled examples with a similar architecture to the previous subtask, only that this time the input to the network will be the hidden representation $\\mathbf{z}=g_\\phi(\\mathbf{x})$ created by the autoencoder in Task 1. Report the accuracy on the test set.\n",
    "3. Compare the results of both models. Which model performed better? Do you have any hypotheses as to why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labeled_data = Subset(train_data, range(100))\n",
    "labeled_loader = torch.utils.data.DataLoader(dataset=labeled_data, batch_size=32, shuffle=True)\n",
    "\n",
    "test_data = MNIST(root='.', \n",
    "                   train=False, \n",
    "                   transform=transforms,    \n",
    "                   download=True)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=5000, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# implement and train the baseline model here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# implement and train the model based on the representation produced by the autoencoder here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[your answer here]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generative models\n",
    "\n",
    "Neural-net based generative models allow us to [generate new faces](https://thispersondoesnotexist.com/) or [generate text](https://transformer.huggingface.co/doc/gpt2-large). The next task will involve creating a generative autoencoder and training it on the FashionMNIST dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms = Compose([ToTensor(), Lambda(lambda x: x.flatten())])\n",
    "train_data = FashionMNIST(root='.', \n",
    "                          train=True, \n",
    "                          transform=transforms,\n",
    "                          download=True)\n",
    "\n",
    "batch_size = 256\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True) \n",
    "\n",
    "fig, axes = plt.subplots(2, 5, figsize=(13, 7))\n",
    "for im, ax in zip(train_data.data[:10], axes.reshape(-1)):\n",
    "    ax.imshow(im, cmap='gray')\n",
    "    ax.set_xticks(())\n",
    "    ax.set_yticks(())\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wasserstein Autoencoder\n",
    "The Wasserstein Autoencoder is identical in architecture to the vanilla one, with the additional constraint that the codes in the latent space are to form a normal distribution. Thanks to this we'll be able to generate new examples by sampling noise from the normal distribution and sending it through the decoder.\n",
    "\n",
    "The loss function is comprised of two parts; the reconstruction loss and a distance between probability distributions:\n",
    "$$L_{WAE-MMD} =\\frac{1}{n}\\sum_i \\lVert\\mathbf{x}_i-f_\\theta(g_\\phi(\\mathbf{x}_i))\\rVert_2^2+C\\cdot \\text{MMD}(g_\\phi(\\mathbf{x}_i),(\\mathbf{y}_j)),$$\n",
    "where $\\mathbf{y}_j$ are samples from the normal distribution $\\mathcal{N}(0, I)$, and $C \\in \\mathbb{R}$ is hyperparameter which weights the different components of the cost function.\n",
    "\n",
    "The formula for Maximum Mean Discrepancy is as follows:\n",
    "$$\\text{MMD}((\\mathbf{y}_i),(\\mathbf{z}_j))=\\frac{1}{n^2}\\sum_{i,i'}k(\\mathbf{y}_i,\\mathbf{y}_{i'})+\\frac{1}{n^2}\\sum_{j,j'}k(\\mathbf{z}_j,\\mathbf{z}_{j'})-\\frac{2}{n^2}\\sum_{i,j}k(\\mathbf{y}_i,\\mathbf{z}_j),$$\n",
    "where $k$ is a kernel function.\n",
    "\n",
    "MMD describes a distance between the hidden representation $\\mathbf{z}=g_\\phi(\\mathbf{x})$, obtained by passing the training examples through the encoder, and samples $\\mathbf{y}_j\\sim\\mathcal{N}(0, I)$. Minimizing this cost will make the distribution produced by the encoder be more like the normal distribution, which is what we want to achieve.\n",
    "\n",
    "We will use the IMQ (inverse multi-quadratic) kernel:\n",
    "$$k(\\mathbf{y}, \\mathbf{z})=\\frac{\\sigma}{\\sigma+\\lvert\\mathbf{y} - \\mathbf{z} \\rvert^2},$$\n",
    "where $\\sigma$ is a hyperparameter you need to find."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3 (1p)\n",
    "Implement the Wasserstein Autoencoder with the Maximum Mean Discrepancy loss component.\n",
    "\n",
    "1. Implement the autoencoder architecture (encoder + decoder) as in Task 1. The architecture should take into account that FashionMNIST is more complicated than MNIST (e.g. use 50 dimensions for the latent space).\n",
    "2. Implement a training loop for WAE, where we minimize the loss function $L_{WAE-MMD}$.\n",
    "3. Find hypeparameters (learning rate, number of training epochs, $C$, $\\sigma$, etc.), so that the reconstruction and generated samples look decent (use the `plot_samples` function below). (Start from $C=1$, $\\sigma=2D$, where $D$ is the dimensionality of the latent space.)\n",
    "\n",
    "Cf. [the paper](https://arxiv.org/abs/1711.01558) and [the official implementation](https://github.com/tolstikhin/wae)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_samples(model):\n",
    "    sampled_z = torch.randn(20, model.latent_dim).to(device)\n",
    "    generated = model.decode(sampled_z)\n",
    "\n",
    "    generated = generated.cpu().detach().numpy()\n",
    "    fig, axes = plt.subplots(2, 10, figsize=(15, 4))\n",
    "    for gen_im, ax in zip(generated, axes.reshape(-1)):\n",
    "        ax.imshow(gen_im.reshape(28, 28), cmap=\"gray\")\n",
    "        ax.set_xticks(())\n",
    "        ax.set_yticks(())\n",
    "    fig.tight_layout()\n",
    "    fig.suptitle(\"Generated samples\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WAEMMD(nn.Module):\n",
    "    \n",
    "    def __init__(self, latent_dim):\n",
    "        \n",
    "        super(WAEMMD, self).__init__()\n",
    "        \n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        self.encoder = ???\n",
    "        self.decoder = ???\n",
    "    \n",
    "    def encode(self, x):\n",
    "        return ???\n",
    "    \n",
    "    def decode(self, encoded):\n",
    "        return ???\n",
    "    \n",
    "    def forward(self, x):\n",
    "        encoded = ???\n",
    "        decoded = ???\n",
    "        return encoded, decoded\n",
    "    \n",
    "    def mmd_loss(self, y, sigma):\n",
    "        ???"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = ???\n",
    "lr = ???\n",
    "latent_dim = ???\n",
    "\n",
    "wae = WAEMMD(latent_dim).to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(wae.parameters(), lr=lr)\n",
    "\n",
    "sigma = ???\n",
    "C = ???\n",
    "\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    epoch_losses = []\n",
    "    for step, (x, _) in enumerate(train_loader):\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        x = x.to(device)\n",
    "        \n",
    "        encoded, decoded = wae(x)\n",
    "\n",
    "        rec_loss = ???\n",
    "        latent_loss = ???\n",
    "        loss = ???\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        epoch_losses += [loss.item()]\n",
    "\n",
    "    print(f'Epoch: {epoch+1} | train loss: {np.mean(epoch_losses):.5f}')\n",
    "\n",
    "    if epoch % 10 == 0:\n",
    "        plot_dataset(train_data, wae)\n",
    "        plot_pca(train_data, wae)\n",
    "        plot_samples(wae)"
   ]
  }
 ],
 "metadata": {
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}