lab4_discrete-autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VPY1-2Q5tCtF"
   },
   "source": [
    "# Lab 4 - (discrete) AutoEncoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l8BHBn76tCtI"
   },
   "source": [
    "## Plan for today\n",
    "\n",
    "* train a basic Autoencoder\n",
    "* try out two ways of differentially sampling from a discrete distribution\n",
    "  * score function\n",
    "  * Gumbel-softmax\n",
    "* train a discrete autoencoder using those two methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kEOcjR0stCtK"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.optim import SGD\n",
    "from torch import nn\n",
    "import torch\n",
    "from typing import List\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision import transforms as tv\n",
    "from torch.utils.data import DataLoader, Subset\n",
    "from sklearn.metrics import accuracy_score\n",
    "from tqdm.notebook import tqdm\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.manifold import TSNE\n",
    "from time import sleep\n",
    "from torchvision.models import vgg16, vgg16_bn, resnet50, resnet18\n",
    "import seaborn as sns\n",
    "from sklearn.cluster import KMeans\n",
    "from collections import Counter\n",
    "from typing import Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "crDkmJH8uPyC",
    "outputId": "a783c960-ad0d-456a-c716-2a3033875227"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 423,
     "referenced_widgets": [
      "1314f6ccf99d49449083da83d1f69d4f",
      "29818f78115d4995ba90b51cfa4c6d09",
      "302509470d454d9cb2b8a237fb53044b",
      "8c3896fc07be4c0b9be84ec1ca23286a",
      "badb80d37cc1461aa8defa3b283a4601",
      "6c2a55007d904910aed7e531ee1105e5",
      "716cf6da026c40c1907a79d4874e6621",
      "8487ca2f999f4a73bdca08cdb9d5ec1c",
      "b9fb44ce45b64171a34ec4cfa17fce6b",
      "0d8d89094fc84b09ab3977e685f0671b",
      "c73ae55b11f94c36bbb6a3781966eb47",
      "67faf8385bd44f6380183689d7797590",
      "1b1946e5822944ea978d7d4e58d1688c",
      "5e35fce4a62346308ca1d07162197b56",
      "122eb176287747afbfe7057771af7083",
      "4ba237fe872e46b8b436f6108569dbd4",
      "7941d24c6996426aa2a69ae1dcc47a16",
      "46ce073e4bc649c5b188796da08f79c4",
      "6d2c2483d76c4759aaec18a27544ff50",
      "0525632d354847c0bc7f7317b46eb8cf",
      "8a0f48916f964f8bb701279393ec106d",
      "1aaa8ed9fec84246b52eac0ed34940f9",
      "3d71d670cfab4cd18051e4ce2df5a157",
      "3a8559e4557f4f7880c0b4c7f53f9988",
      "1165c611fe16436985d7a38bee1abd4c",
      "dd989555ec1a410d9491cbef37ba81b3",
      "a80d958f1abd45f6938ae74910e2e72b",
      "fc6eda47063245b4bb8d57059944cbe5",
      "87139e76d6744794a9933be9d5688996",
      "f6533848b8b1451999cfabb3bafc1b36",
      "0a83ae6c2d214eeaa7688b4215bf9503",
      "d74babbad33944b6ae5cae05b8cff5b3",
      "b7b2d02e68354a2e928e56758acf3c0d",
      "beb7c81b283a4456bd9c5470586ba7d0",
      "689043476e8142e8887345cffac50b6e",
      "32c51d14bf10485aa75c44822c0cad40",
      "bbb8f55bb25241b4992962e2f33d357d",
      "93ce50324c7849088cead4507d613007",
      "bca6747559e34f66bdf49e329be7a7b5",
      "2b10578963d04ebfb6038db2f18f785b",
      "2d2877e679e74283b22e4f1502e17464",
      "36039b3cb57c4e68964c24d7d27b20e9",
      "6c7319b57a1a43028650c34e358134ea",
      "ef7fb768cd194f07b26c0b5fbf5468b8"
     ]
    },
    "id": "VmRBEEmHue3G",
    "outputId": "5c2cc002-91bd-4ec6-9269-f0c490365ccc"
   },
   "outputs": [],
   "source": [
    "\n",
    "ds_train = MNIST(root=\"data\", train=True, download=True, transform=tv.ToTensor())\n",
    "\n",
    "ds_test = MNIST(root=\"data\", train=False, download=True, transform=tv.ToTensor())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qwT8wAyVwHyj"
   },
   "outputs": [],
   "source": [
    "batch_size=128\n",
    "dl_train = DataLoader(ds_train, batch_size, shuffle=True, drop_last=False) # dataloader with full dataset \n",
    "\n",
    "dl_test = DataLoader(ds_test, batch_size, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wPQNkFmYh6NO"
   },
   "source": [
    "## 1 - Warmup: vanilla AutoEncoder\n",
    "\n",
    "Let's start with training a classic AutoEncoder on MNIST. The architecture and training loop have been provided below.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NCEp1Kcnk5D-"
   },
   "outputs": [],
   "source": [
    "class Reshape(nn.Module):\n",
    "  def __init__(self, shape):\n",
    "    super().__init__()\n",
    "    self.shape = shape\n",
    "  \n",
    "  def forward(self, x):\n",
    "    return x.reshape(-1, *self.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vbSHwkTwtCtO"
   },
   "outputs": [],
   "source": [
    "class AE(nn.Module):\n",
    "  def __init__(self, hidden_sizes: List[int], downsize_steps: List[int], chw: int = (1, 28, 28), latent_size: int = 128):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "      hidden sizes - hidden size of each consecutive convolution\n",
    "      downsize_steps - numbers of convolutions for which there will be stride = 2\n",
    "      chw - size of the input image\n",
    "      latent_size: latent dimension\n",
    "    \n",
    "    Forward method returns the reconstructed batch *and* the latent vectors.\n",
    "    \"\"\"\n",
    "    super().__init__()\n",
    "\n",
    "    c, h, w = chw\n",
    "\n",
    "    conv_1 = nn.Conv2d(c, hidden_sizes[0], 3, padding=1)\n",
    "    conv_n = nn.ConvTranspose2d(hidden_sizes[0], c, 3, padding=1)\n",
    "\n",
    "    encoder = []\n",
    "    decoder = []\n",
    "\n",
    "    ht, wt = h, w\n",
    "\n",
    "    for i, (in_hs, out_hs) in enumerate(zip(hidden_sizes[:-1], hidden_sizes[1:])):\n",
    "      downsize = (i in downsize_steps)\n",
    "      encoder.append(nn.Sequential(\n",
    "          nn.Conv2d(\n",
    "              in_hs, out_hs, 3, padding=1,\n",
    "              stride=2 if downsize else 1\n",
    "          ),\n",
    "          nn.BatchNorm2d(out_hs),\n",
    "          nn.ReLU()\n",
    "      ))\n",
    "\n",
    "      decoder = [nn.Sequential(\n",
    "          nn.ConvTranspose2d( \n",
    "              out_hs, in_hs, 3, padding=1, output_padding=1,\n",
    "              stride=2 if downsize else 1,\n",
    "          ),\n",
    "          nn.BatchNorm2d(in_hs),\n",
    "          nn.ReLU()\n",
    "      )] + decoder\n",
    "      if downsize:\n",
    "        ht = ht // 2\n",
    "        wt = wt // 2\n",
    "    \n",
    "    le = nn.Sequential(\n",
    "        nn.Flatten(),\n",
    "        nn.Linear(ht * wt * hidden_sizes[-1], latent_size)\n",
    "    )\n",
    "\n",
    "    ld = nn.Sequential(\n",
    "        nn.Linear(latent_size, ht * wt * hidden_sizes[-1]),\n",
    "        Reshape((  hidden_sizes[-1], ht, wt))\n",
    "    )\n",
    "    self.encoder = nn.Sequential(conv_1, *encoder, le)\n",
    "    self.decoder = nn.Sequential(ld, *decoder, conv_n)\n",
    "\n",
    "\n",
    "  def forward(self, x):\n",
    "    l = self.encoder(x)\n",
    "    d = self.decoder(l)\n",
    "    return d, l\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4TMHBMmxswo0"
   },
   "outputs": [],
   "source": [
    "def draw_reconstruction(original: torch.Tensor, expected: torch.Tensor, epoch: int):\n",
    "  assert len(original) == len(expected)\n",
    "  fig, ax = plt.subplots(nrows=2, ncols=len(original), figsize=(3*len(original), 6))\n",
    "  \n",
    "  for i, (o,e) in enumerate(zip(original, expected)):\n",
    "    ax[0][i].imshow(o.squeeze())\n",
    "    ax[1][i].imshow(e.squeeze())\n",
    "  plt.suptitle(f\"{epoch}: Original / expected\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VNAMQigAtCtQ"
   },
   "outputs": [],
   "source": [
    "def train_ae(net, train_dl, test_dl, num_epochs: int = 20, lr: float = 3e-4, loss_fn = None):\n",
    "\n",
    "  net = net.to(device)\n",
    "  loss_fn = loss_fn or nn.MSELoss()\n",
    "  opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=3e-4)  \n",
    "  \n",
    "  history = []\n",
    "\n",
    "  epoch_progress = tqdm(list(range(num_epochs)))\n",
    "  for i in epoch_progress:\n",
    "      train_loss = 0\n",
    "      net.train()\n",
    "      for iteration, (X_train, y_train) in enumerate(train_dl):\n",
    "          X_train, y_train = [t.to(device) for t in [X_train, y_train]]\n",
    "          opt.zero_grad()\n",
    "          latent = net.encoder(X_train)\n",
    "          X_pred = net.decoder(latent)\n",
    "          loss = loss_fn(X_train, X_pred)\n",
    "          loss.backward()\n",
    "          opt.step()\n",
    "          train_loss += loss.item()\n",
    "          \n",
    "      val_loss = 0\n",
    "      net.eval()\n",
    "      with torch.no_grad():\n",
    "          for iteration, (X_val, y_val) in enumerate(test_dl):\n",
    "              X_val, y_val = [t.to(device) for t in [X_val, y_val]]\n",
    "\n",
    "              X_pred, _ = net(X_val)\n",
    "              if iteration == 0:\n",
    "                draw_reconstruction(\n",
    "                    X_val[:10].detach().cpu().numpy(), X_pred[:10].detach().cpu().numpy(), epoch=i\n",
    "                )\n",
    "              loss = loss_fn(X_val, X_pred)\n",
    "              val_loss += loss.item()\n",
    "      \n",
    "      epoch_progress.set_description(f'#Epoch: {i}, train loss: {train_loss:.2f}, val loss: {val_loss:.2f}')\n",
    "      history.append({\"e\": i, \"train_loss\": train_loss, \"val_loss\": val_loss})\n",
    "    \n",
    "  for m in [\"train_loss\", \"val_loss\"]:\n",
    "      plt.plot([h[\"e\"] for h in history], [h[m] for h in history], label=m)\n",
    "  \n",
    "  plt.grid()\n",
    "  plt.legend()\n",
    "  plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CbKYRO6CkWYa"
   },
   "source": [
    "\n",
    "### Task for you - train the autoencoder\n",
    "\n",
    "* a relatively small model (2 convolutions in the encoder and 2 convolutions in the decoder) should suffice\n",
    "* try out 2 different latent sizes \n",
    "* try out [MSELoss](https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html) or [BCEWithLogitsLoss](https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L3sb-xRxjNWu"
   },
   "outputs": [],
   "source": [
    "ae = AE(..., latent_size=...)\n",
    "\n",
    "train_ae(ae, train_dl=dl_train, test_dl=dl_test, num_epochs=10, loss_fn=...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hRE7YWlNk-Lj"
   },
   "source": [
    "**Question 1** - What are the advantages / disadvantages of using MSE and BCE as reconstruction losses? Which one, in your opinion, makes more sense for this dataset? Which one would you use for a dataset such as CIFAR-10?\n",
    "\n",
    "**YOUR ANSWER HERE**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RRXPz2kokhsG"
   },
   "source": [
    "### Task for you - visualize the latent space\n",
    "\n",
    "Run the autoencoder on the validation dataset and visualize the latent embeddings on 2-D space using T-SNE or PCA. Please color each embedding according to it's ground-truth label.\n",
    "\n",
    "Since the AE returns a tuple of `<reconstructed image>, <latent vector>`, you should be able to re-use the `draw_embeddings` function from Lab 3. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 352
    },
    "id": "wILP0ioIzlMk",
    "outputId": "354d9d9f-1bad-4f53-c485-3651559f838b"
   },
   "outputs": [],
   "source": [
    "def draw_embeddings(net, dl):\n",
    "  # YOUR CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8ajSgLEimT6U"
   },
   "source": [
    "## 2 - Discrete AutoEncoder\n",
    "\n",
    "We would like our AutoEncoder to model a **discrete** data distribution. \n",
    "\n",
    "In other words, if the latent of the AE vector has shape $K$, we want it to take only $K$ different values, represented by one-hot vectors. Notice that this is equivalent to **clustering** - the encoder assigns $x$ into one of the $K$ clusters!\n",
    "\n",
    "\n",
    "How to achieve that in practice?\n",
    "* the Encoder should transform each image $x$ into a one-hot vector $z$ of shape $K$\n",
    "  * we can apply e.g. $softmax$ to the latent vector output by the Encoder in order to obtain the parameters of discrete data distribution $p$\n",
    "* the Decoder should transform a one-hot vector of shape $K$ sampled from $p$ into something as similar to $x$ as possible.\n",
    "* we should minimize the **expected** value of cost: \n",
    "\n",
    "$$\n",
    "\\mathbb{E}_{z \\sim softmax(E(x))} || x - D(z)||^2\n",
    "$$\n",
    "\n",
    "\n",
    "The rough training procedure for a sample $x$ will be to:\n",
    "1. Estimate the distribution: $ p(x) = softmax(E(x)) = [\\pi_1, ..., \\pi_K]$ with the encoder\n",
    "2. Sample $M$ one-hot vectors $(z_1, ..., z_M)$ from distribution $p(x)$\n",
    "3. Use the Decoder to transform $(z_1, ..., z_M)$ into $(\\hat{x}_1, ..., \\hat{x}_M)$\n",
    "4. Minimize the cost: $\\frac{1}{M} \\sum_{i=1}^M ||x - \\hat{x_i}||^2$\n",
    "\n",
    "Step **2** is especially tricky - *how to sample from the distribution generated by the encoder in a differentiable way*? \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "__GdwLAzso6h"
   },
   "source": [
    "### 2a - score function\n",
    "\n",
    "One of the ways we can do this is to *inject* $p(x)$ into the gradient between the encoder and decoder:\n",
    "\n",
    "For an image $x$ our Encoder generates distribution parameters $p(x)$ - a vector of shape $K$. We then sample $M$ one-hot vectors from $p(x)$ - we will denote them as $S$ - a matrix of size $(M, K)$, i.e. ($S \\in \\mathbb{R}^{M \\times K}$). $S$ serves as input to the decoder which decodes it into $M$ predicted images $\\hat{x}_{1...M}$.\n",
    "\n",
    "The loss will be formulated as above: $\\frac{1}{M} \\sum_{i=1}^M ||x - \\hat{x_i}||^2$ - we will calculate the gradients of the **decoder** with respect to this loss, just like in an ordinary Autoencoder (in torch we do this with the `.backward()` function, which calculates the gradients, but does not apply them). \n",
    "\n",
    "This operation will give us gradients of:\n",
    "* decoder's parameters\n",
    "* sampled decoder inputs - $\\nabla S \\in \\mathbb{R}^{M \\times K}$\n",
    "\n",
    "The `backward()` procedure will not be able to go further, as decoder inputs ($S$) are the results of sampling!\n",
    "\n",
    "In order to propagate $\\nabla S$ to the Encoder, we will calculate a matrix of size $(M,K)$: \n",
    "$$\\nabla S \\cdot S \\cdot log(p(x))$$ \n",
    "and average it along the first axis ($M$ - i.e. the number of our samples from $p(x)$). This results in a vector of shape $K$ - i.e. the same shape as $p(x)$.\n",
    "\n",
    "Finally, we can pass the constructed vector into the `encoder.backward(...)` function, thus calculating encoder gradients.\n",
    "Note that all of the above was conducted for a **single** sample $x$. When we will work with batches of samples, we need to remember to take into account the batch dimension as well!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "__GdwLAzso6h"
   },
   "source": [
    "The above method is called score function. In this case, the operation of sampling is non-differentiable, but the gradient flows from decoder to the encoder through the values of vector $p(x) = softmax(E(x))$.\n",
    "\n",
    "In a case where $K >> M$, most of the elements of \"gradient\" passed to the encoder will be zeroed-out, which causes a need for large $M$ for the scoring function to work well - this is one of the drawbacks of this method.\n",
    "\n",
    "\n",
    "### Task for you - train a discrete VAE using scoring function\n",
    "* modify the below training loop to calculate the loss function \n",
    "* use a latent size bigger than number of classes, e.g. 64\n",
    "* try out a couple of different values of $M$, e.g. 1, 2, 4.\n",
    "* in each validation step you can simply transform the latent vector into one-hot by setting $1$ in the postition of the largest element of $p$\n",
    "* after the training, draw the embeddings generated by the encoder. How many clusters are there?\n",
    "\n",
    "**As you have seen, the scoring function is very much imperfect - if you have trouble training the autoencoder with this method, please don't spend too much time on it and proceed to task 2B!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ut5oaQANgs05"
   },
   "outputs": [],
   "source": [
    "def train_discrete_ae_gumbel(net, train_dl, test_dl, num_epochs: int = 20, lr: float = 3e-4, loss_fn = None):\n",
    "\n",
    "  net = net.to(device)\n",
    "  loss_fn = loss_fn or nn.MSELoss()\n",
    "  opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=3e-4)  \n",
    "  \n",
    "  history = []\n",
    "\n",
    "  epoch_progress = tqdm(list(range(num_epochs)))\n",
    "  for i in epoch_progress:\n",
    "      train_loss = 0\n",
    "      net.train()\n",
    "      for iteration, (X_train, y_train) in enumerate(train_dl):\n",
    "          X_train, y_train = [t.to(device) for t in [X_train, y_train]]\n",
    "          opt.zero_grad()\n",
    "          \n",
    "          p = net.encoder(X_train)\n",
    "          ###### YOUR CODE HERE ########\n",
    "          # * sample one-hots from p (those tensors should have set requires_grad to True!)\n",
    "          # * pass one-hots through decoder \n",
    "          # * calculate decoder loss\n",
    "          # * calculate decoder gradient\n",
    "          # * do one-hots also have the gradient calculated?\n",
    "          # * calculate the \"gradient\" for the encoder, as described above\n",
    "          # * pass the \"gradient\" through encoder with encoder.backward(my_calculated_gradient)\n",
    "          \n",
    "          ##############################\n",
    "\n",
    "          opt.step()\n",
    "          train_loss += loss.item()\n",
    "          \n",
    "      val_loss = 0\n",
    "      net.eval()\n",
    "      with torch.no_grad():\n",
    "          for iteration, (X_val, y_val) in enumerate(test_dl):\n",
    "              X_val, y_val = [t.to(device) for t in [X_val, y_val]]\n",
    "\n",
    "              latent = net.encoder(X_val)\n",
    "              softmax_latent = torch.softmax(latent, -1)\n",
    "\n",
    "              ### YOUR CODE here\n",
    "              # transform the softmax_latent to one-hot vectors \n",
    "              onehot_latent = ...\n",
    "              \n",
    "              ###\n",
    "              X_pred = net.decoder(onehot_latent)\n",
    "              if iteration == 0:\n",
    "                draw_reconstruction(\n",
    "                    X_val[:10].detach().cpu().numpy(), X_pred[:10].detach().cpu().numpy(), epoch=i\n",
    "                )\n",
    "              loss = loss_fn(X_val, X_pred)\n",
    "              val_loss += loss.item()\n",
    "      \n",
    "      epoch_progress.set_description(f'#Epoch: {i}, train loss: {train_loss:.2f}, val loss: {val_loss:.2f}')\n",
    "      history.append({\"e\": i, \"train_loss\": train_loss, \"val_loss\": val_loss})\n",
    "    \n",
    "  for m in [\"train_loss\", \"val_loss\"]:\n",
    "      plt.plot([h[\"e\"] for h in history], [h[m] for h in history], label=m)\n",
    "  \n",
    "  plt.grid()\n",
    "  plt.legend()\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lEPdCivuV_wU"
   },
   "source": [
    "### 2b - Gumbel-softmax\n",
    "\n",
    "An alternative way is to approximate the discrete distribution with a continuous one, through a Gumbel-softmax trick.\n",
    "\n",
    "Suppose we have a discrete distribution with probabilities:\n",
    "\n",
    "$$\\pi = [\\pi_1, ..., \\pi_K]$$\n",
    "\n",
    "In order to sample from such a distribution, we can sample a vector of $K$ values from a uniform distribution:\n",
    "\n",
    "$$\n",
    "u = [u_1, ..., u_K];\\\\\n",
    "u \\sim Unif(0,1)\\\n",
    "$$\n",
    "\n",
    "then, transform it into a vector $g$:\n",
    "\n",
    "$$g = -log(-log(u)) = [g_1, ..., g_K]$$\n",
    "\n",
    "Let's now sum $g$ and $log(\\pi)$ and divide it by **temperature** $\\tau$. Finally, let's apply $softmax$ to the resulting vector:\n",
    "\n",
    "$$\n",
    "y = softmax((g + log(\\pi)) / \\tau)\n",
    "$$\n",
    "\n",
    "The result of this operation is a vector of $K$ values which sums up to 1. Since vectors $g$ come from a continuous uniform distribution and vectors $y$ are a function of $g$, $\\pi$ and $\\tau$, it turns out that vectors $y$ come from a continuous distribution with parameters $\\pi$ and $\\tau$.\n",
    "\n",
    "### Task for you - please implement sampling from such a distribution\n",
    "\n",
    "* select a non-trivial categorical distribution $\\pi$\n",
    "* try out sampling for different values of $\\tau$, ranging between 10^2 ... 10^-5\n",
    "* for each $\\tau$ sample 1000 vectors and calculate the mean sample\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5Le9CLvpOqra",
    "outputId": "b7c0072d-10b9-46af-c5aa-e9a30df67b7a"
   },
   "outputs": [],
   "source": [
    "def gumbel_sample(p: torch.Tensor, t: float, n: int):\n",
    "  \"\"\"\n",
    "  p: tensor of shape [d] - represents a categorical distribution with `d` categories - i.e. p should consists of d non-negative values which sum to 1.\n",
    "  t: temperature\n",
    "  n: number of samples\n",
    "\n",
    "  returns:\n",
    "    a tensor of shape [n, d] with n samples from distribution Y\n",
    "  \"\"\"\n",
    "\n",
    "\n",
    "\n",
    "p = torch.tensor([.25, .50, .15, .1]).float()\n",
    "\n",
    "gumbel_sample(p, t=100, n=1000).mean(dim=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U-TtrzuI32z1"
   },
   "source": [
    "**Question 2** - what have you noticed for large and small values of $\\tau$?\n",
    "\n",
    "**YOUR ANSWER HERE**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "25c-OxuP4New"
   },
   "source": [
    "### Task for you - train a discrete autoencoder with Gumbel-softmax\n",
    "\n",
    "* in each training step, for each image $x$:\n",
    "  * generate distribution $E(x) = \\pi_x$ with the decoder\n",
    "  * sample $M$ one-hot vectors $[z_1^x, ..., z_M^x]$ from $\\pi_x$\n",
    "  * use the decoder to transform $[z_1^x, ..., z_M^x]$ into $[\\hat{x_1}, ..., \\hat{x_M}]$\n",
    "  * minimize the mean difference (MSE or BCE) between $x$ and $[\\hat{x_1}, ..., \\hat{x_M}]$  \n",
    "* with each epoch, decrease the temperature $\\tau$\n",
    "* in each validation step you can simply transform the latent vector into one-hot using softmax\n",
    "\n",
    "After training:\n",
    "* visualize the latent clusters with colors from ground-truth labels (just like with the vanilla AE)\n",
    "* for each dimension of the latent generate a one-hot vector and visualize the output of the decoder for that vector. For example, if our latent size was $4$, we would visualize in total four decoder outputs for the following one-hot vectors:\n",
    "  * $[1, 0, 0, 0]$\n",
    "  * $[0, 1, 0, 0]$\n",
    "  * $[0, 0, 1, 0]$\n",
    "  * $[0, 0, 0, 1]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7Cg6XLuCW3Ne"
   },
   "outputs": [],
   "source": [
    "def train_discrete_ae_gumbel(net, train_dl, test_dl, num_epochs: int = 20, lr: float = 3e-4, loss_fn = None):\n",
    "\n",
    "  net = net.to(device)\n",
    "  loss_fn = loss_fn or nn.MSELoss()\n",
    "  opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=3e-4)  \n",
    "  \n",
    "  history = []\n",
    "\n",
    "  epoch_progress = tqdm(list(range(num_epochs)))\n",
    "  for i in epoch_progress:\n",
    "      train_loss = 0\n",
    "      net.train()\n",
    "      for iteration, (X_train, y_train) in enumerate(train_dl):\n",
    "          X_train, y_train = [t.to(device) for t in [X_train, y_train]]\n",
    "          opt.zero_grad()\n",
    "          \n",
    "          latent = net.encoder(X_train)\n",
    "          ###### YOUR CODE HERE ########\n",
    "\n",
    "          loss = ...\n",
    "          ##############################\n",
    "\n",
    "          loss.backward()\n",
    "          opt.step()\n",
    "          train_loss += loss.item()\n",
    "          \n",
    "      val_loss = 0\n",
    "      net.eval()\n",
    "      with torch.no_grad():\n",
    "          for iteration, (X_val, y_val) in enumerate(test_dl):\n",
    "              X_val, y_val = [t.to(device) for t in [X_val, y_val]]\n",
    "\n",
    "              latent = net.encoder(X_val)\n",
    "              softmax_latent = torch.softmax(latent, -1)\n",
    "\n",
    "              ### YOUR CODE here\n",
    "              # transform the softmax_latent to one-hot vectors \n",
    "              onehot_latent = ...\n",
    "              \n",
    "              ###\n",
    "              X_pred = net.decoder(onehot_latent)\n",
    "              if iteration == 0:\n",
    "                draw_reconstruction(\n",
    "                    X_val[:10].detach().cpu().numpy(), X_pred[:10].detach().cpu().numpy(), epoch=i\n",
    "                )\n",
    "              loss = loss_fn(X_val, X_pred)\n",
    "              val_loss += loss.item()\n",
    "      \n",
    "      epoch_progress.set_description(f'#Epoch: {i}, train loss: {train_loss:.2f}, val loss: {val_loss:.2f}')\n",
    "      history.append({\"e\": i, \"train_loss\": train_loss, \"val_loss\": val_loss})\n",
    "    \n",
    "  for m in [\"train_loss\", \"val_loss\"]:\n",
    "      plt.plot([h[\"e\"] for h in history], [h[m] for h in history], label=m)\n",
    "  \n",
    "  plt.grid()\n",
    "  plt.legend()\n",
    "  plt.show()\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "wzum_lab4.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "uj",
   "language": "python",
   "name": "uj"
  },
  "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.16"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0525632d354847c0bc7f7317b46eb8cf": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "0a83ae6c2d214eeaa7688b4215bf9503": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "0d8d89094fc84b09ab3977e685f0671b": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1165c611fe16436985d7a38bee1abd4c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_f6533848b8b1451999cfabb3bafc1b36",
      "max": 1648877,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_0a83ae6c2d214eeaa7688b4215bf9503",
      "value": 1648877
     }
    },
    "122eb176287747afbfe7057771af7083": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_8a0f48916f964f8bb701279393ec106d",
      "placeholder": "​",
      "style": "IPY_MODEL_1aaa8ed9fec84246b52eac0ed34940f9",
      "value": " 29696/? [00:00&lt;00:00, 6844.41it/s]"
     }
    },
    "1314f6ccf99d49449083da83d1f69d4f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_29818f78115d4995ba90b51cfa4c6d09",
       "IPY_MODEL_302509470d454d9cb2b8a237fb53044b",
       "IPY_MODEL_8c3896fc07be4c0b9be84ec1ca23286a"
      ],
      "layout": "IPY_MODEL_badb80d37cc1461aa8defa3b283a4601"
     }
    },
    "1aaa8ed9fec84246b52eac0ed34940f9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "1b1946e5822944ea978d7d4e58d1688c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7941d24c6996426aa2a69ae1dcc47a16",
      "placeholder": "​",
      "style": "IPY_MODEL_46ce073e4bc649c5b188796da08f79c4",
      "value": ""
     }
    },
    "29818f78115d4995ba90b51cfa4c6d09": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6c2a55007d904910aed7e531ee1105e5",
      "placeholder": "​",
      "style": "IPY_MODEL_716cf6da026c40c1907a79d4874e6621",
      "value": ""
     }
    },
    "2b10578963d04ebfb6038db2f18f785b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "2d2877e679e74283b22e4f1502e17464": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "302509470d454d9cb2b8a237fb53044b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_8487ca2f999f4a73bdca08cdb9d5ec1c",
      "max": 9912422,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_b9fb44ce45b64171a34ec4cfa17fce6b",
      "value": 9912422
     }
    },
    "32c51d14bf10485aa75c44822c0cad40": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_2d2877e679e74283b22e4f1502e17464",
      "max": 4542,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_36039b3cb57c4e68964c24d7d27b20e9",
      "value": 4542
     }
    },
    "36039b3cb57c4e68964c24d7d27b20e9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "3a8559e4557f4f7880c0b4c7f53f9988": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_fc6eda47063245b4bb8d57059944cbe5",
      "placeholder": "​",
      "style": "IPY_MODEL_87139e76d6744794a9933be9d5688996",
      "value": ""
     }
    },
    "3d71d670cfab4cd18051e4ce2df5a157": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_3a8559e4557f4f7880c0b4c7f53f9988",
       "IPY_MODEL_1165c611fe16436985d7a38bee1abd4c",
       "IPY_MODEL_dd989555ec1a410d9491cbef37ba81b3"
      ],
      "layout": "IPY_MODEL_a80d958f1abd45f6938ae74910e2e72b"
     }
    },
    "46ce073e4bc649c5b188796da08f79c4": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "4ba237fe872e46b8b436f6108569dbd4": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5e35fce4a62346308ca1d07162197b56": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6d2c2483d76c4759aaec18a27544ff50",
      "max": 28881,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_0525632d354847c0bc7f7317b46eb8cf",
      "value": 28881
     }
    },
    "67faf8385bd44f6380183689d7797590": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_1b1946e5822944ea978d7d4e58d1688c",
       "IPY_MODEL_5e35fce4a62346308ca1d07162197b56",
       "IPY_MODEL_122eb176287747afbfe7057771af7083"
      ],
      "layout": "IPY_MODEL_4ba237fe872e46b8b436f6108569dbd4"
     }
    },
    "689043476e8142e8887345cffac50b6e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_bca6747559e34f66bdf49e329be7a7b5",
      "placeholder": "​",
      "style": "IPY_MODEL_2b10578963d04ebfb6038db2f18f785b",
      "value": ""
     }
    },
    "6c2a55007d904910aed7e531ee1105e5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6c7319b57a1a43028650c34e358134ea": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6d2c2483d76c4759aaec18a27544ff50": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "716cf6da026c40c1907a79d4874e6621": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "7941d24c6996426aa2a69ae1dcc47a16": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8487ca2f999f4a73bdca08cdb9d5ec1c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "87139e76d6744794a9933be9d5688996": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "8a0f48916f964f8bb701279393ec106d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8c3896fc07be4c0b9be84ec1ca23286a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_0d8d89094fc84b09ab3977e685f0671b",
      "placeholder": "​",
      "style": "IPY_MODEL_c73ae55b11f94c36bbb6a3781966eb47",
      "value": " 9913344/? [00:00&lt;00:00, 20122875.83it/s]"
     }
    },
    "93ce50324c7849088cead4507d613007": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a80d958f1abd45f6938ae74910e2e72b": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b7b2d02e68354a2e928e56758acf3c0d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "b9fb44ce45b64171a34ec4cfa17fce6b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "badb80d37cc1461aa8defa3b283a4601": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "bbb8f55bb25241b4992962e2f33d357d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6c7319b57a1a43028650c34e358134ea",
      "placeholder": "​",
      "style": "IPY_MODEL_ef7fb768cd194f07b26c0b5fbf5468b8",
      "value": " 5120/? [00:00&lt;00:00, 9818.93it/s]"
     }
    },
    "bca6747559e34f66bdf49e329be7a7b5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "beb7c81b283a4456bd9c5470586ba7d0": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_689043476e8142e8887345cffac50b6e",
       "IPY_MODEL_32c51d14bf10485aa75c44822c0cad40",
       "IPY_MODEL_bbb8f55bb25241b4992962e2f33d357d"
      ],
      "layout": "IPY_MODEL_93ce50324c7849088cead4507d613007"
     }
    },
    "c73ae55b11f94c36bbb6a3781966eb47": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "d74babbad33944b6ae5cae05b8cff5b3": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "dd989555ec1a410d9491cbef37ba81b3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d74babbad33944b6ae5cae05b8cff5b3",
      "placeholder": "​",
      "style": "IPY_MODEL_b7b2d02e68354a2e928e56758acf3c0d",
      "value": " 1649664/? [00:00&lt;00:00, 3285181.71it/s]"
     }
    },
    "ef7fb768cd194f07b26c0b5fbf5468b8": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "f6533848b8b1451999cfabb3bafc1b36": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "fc6eda47063245b4bb8d57059944cbe5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}