homeworks/hw4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mfuMscdgelAu"
   },
   "source": [
    "## Task 1: Theory (3pt)\n",
    "\n",
    "In lecture 7 we proved the [ELBO surgery](http://approximateinference.org/accepted/HoffmanJohnson2016.pdf) theorem:\n",
    "$$\n",
    "    \\frac{1}{n} \\sum_{i=1}^n KL(q(\\mathbf{z} | \\mathbf{x}_i) || p(\\mathbf{z})) = KL(q(\\mathbf{z}) || p(\\mathbf{z})) + \\mathbb{I}_{q} [\\mathbf{x}, \\mathbf{z}],\n",
    "$$\n",
    "where the first term is $KL(q(\\mathbf{z}) || p(\\mathbf{z}))$ includes the aggregated posterior distribution $q(\\mathbf{z})$ and the prior distribution $p(\\mathbf{z})$. Our goal now is to deal with the second term. At the lecture, the second term was equal to:\n",
    "\n",
    "$$\n",
    "    \\mathbb{I}_{q} [\\mathbf{x}, \\mathbf{z}] = \\frac{1}{n}\\sum_{i=1}^n KL(q(\\mathbf{z} | \\mathbf{x}_i) || q (\\mathbf{z})).\n",
    "$$\n",
    "In fact, this is a mutual information between $\\mathbf{x}$ and $\\mathbf{z}$ on the empirical distribution of data and the distribution of $q(\\mathbf{z} | \\mathbf{x})$. Let treat the index of the sample $i$ as a random variable.\n",
    "$$\n",
    "    q(i, \\mathbf{z}) = q(i) q(\\mathbf{z} | i); \\quad p(i, \\mathbf{z}) = p(i) p(\\mathbf{z}); \\quad \n",
    "    q(i) = p(i) = \\frac{1}{n}.\n",
    "$$\n",
    "$$\n",
    "    \\quad q(\\mathbf{z} | i) = q(\\mathbf{z} | \\mathbf{x}_i) \\quad q(\\mathbf{z}) = \\sum_{i=1}^n q(i, \\mathbf{z}) = \\frac{1}{n} \\sum_{i=1}^n q(\\mathbf{z} | \\mathbf{x}_i);  \n",
    "$$\n",
    "Mutual information is a measure of independence between two random variables.\n",
    "$$\n",
    "\t\\mathbb{I}_{q} [\\mathbf{x}, \\mathbf{z}] = \\mathbb{E}_{q(i, \\mathbf{z})} \\log \\frac{q(i, \\mathbf{z})}{q(i)q(\\mathbf{z})}.\n",
    "$$\n",
    "Prove that 2 expressions for mutual information are equal to each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ieenT6NCp_OK"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from collections import defaultdict\n",
    "from sklearn.datasets import make_moons\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "from torchvision.utils import make_grid\n",
    "from torch.distributions.uniform import Uniform\n",
    "from torch.distributions.normal import Normal\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jQeQeajnkCFV",
    "outputId": "90022d5d-9f33-451a-bd9a-7acbc5d9f47f"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MFQTGwLhUOsa"
   },
   "source": [
    "These functions are the helpers that will train your models and visualize the results. You do not have to change them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VAP9CQYsbldo"
   },
   "outputs": [],
   "source": [
    "def train_epoch(model, train_loader, optimizer, use_cuda, loss_key='total'):\n",
    "    model.train()\n",
    "  \n",
    "    stats = defaultdict(list)\n",
    "    for x in train_loader:\n",
    "        if use_cuda:\n",
    "            x = x.cuda()\n",
    "        losses = model.loss(x)\n",
    "        optimizer.zero_grad()\n",
    "        losses[loss_key].backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        for k, v in losses.items():\n",
    "            stats[k].append(v.item())\n",
    "    return stats\n",
    "\n",
    "\n",
    "def eval_model(model, data_loader, use_cuda):\n",
    "    model.eval()\n",
    "    stats = defaultdict(float)\n",
    "    with torch.no_grad():\n",
    "        for x in data_loader:\n",
    "            if use_cuda:\n",
    "                x = x.cuda()\n",
    "            losses = model.loss(x)\n",
    "            for k, v in losses.items():\n",
    "                stats[k] += v.item() * x.shape[0]\n",
    "\n",
    "        for k in stats.keys():\n",
    "            stats[k] /= len(data_loader.dataset)\n",
    "    return stats\n",
    "\n",
    "\n",
    "def train_model(model, train_loader, test_loader, epochs, lr, use_tqdm=False, use_cuda=False, loss_key='total_loss'):\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "\n",
    "    train_losses = defaultdict(list)\n",
    "    test_losses = defaultdict(list)\n",
    "    forrange = tqdm(range(epochs)) if use_tqdm else range(epochs)\n",
    "    if use_cuda:\n",
    "        model = model.cuda()\n",
    "    for epoch in forrange:\n",
    "        model.train()\n",
    "        train_loss = train_epoch(model, train_loader, optimizer, use_cuda, loss_key)\n",
    "        test_loss = eval_model(model, test_loader, use_cuda)\n",
    "\n",
    "        for k in train_loss.keys():\n",
    "            train_losses[k].extend(train_loss[k])\n",
    "            test_losses[k].append(test_loss[k])\n",
    "    return dict(train_losses), dict(test_losses)\n",
    "\n",
    "\n",
    "def plot_training_curves(train_losses, test_losses):\n",
    "    n_train = len(train_losses[list(train_losses.keys())[0]])\n",
    "    n_test = len(test_losses[list(train_losses.keys())[0]])\n",
    "    x_train = np.linspace(0, n_test - 1, n_train)\n",
    "    x_test = np.arange(n_test)\n",
    "\n",
    "    plt.figure()\n",
    "    for key, value in train_losses.items():\n",
    "        plt.plot(x_train, value, label=key + '_train')\n",
    "\n",
    "    for key, value in test_losses.items():\n",
    "        plt.plot(x_test, value, label=key + '_test')\n",
    "\n",
    "    plt.legend()\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def load_pickle(path, flatten=False, binarize=False):\n",
    "    with open(path, 'rb') as f:\n",
    "        data = pickle.load(f)\n",
    "    train_data = data['train'].astype('float32')\n",
    "    test_data = data['test'].astype('float32')\n",
    "    if binarize:\n",
    "        train_data = (train_data > 128).astype('float32')\n",
    "        test_data = (test_data > 128).astype('float32')\n",
    "    else:\n",
    "        train_data = train_data / 255.\n",
    "        test_data = test_data / 255.\n",
    "    train_data = np.transpose(train_data, (0, 3, 1, 2))\n",
    "    test_data = np.transpose(test_data, (0, 3, 1, 2))\n",
    "    if flatten:\n",
    "        train_data = train_data.reshape(len(train_data.shape[0]), -1)\n",
    "        test_data = test_data.reshape(len(train_data.shape[0]), -1)\n",
    "    return train_data, test_data\n",
    "\n",
    "\n",
    "def show_samples(samples, title, nrow=10):\n",
    "    samples = torch.FloatTensor(samples)\n",
    "    grid_img = make_grid(samples, nrow=nrow)\n",
    "    plt.figure()\n",
    "    plt.title(title)\n",
    "    plt.imshow(grid_img.permute(1, 2, 0))\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def visualize_data(data, title):\n",
    "    idxs = np.random.choice(len(data), replace=False, size=(100,))\n",
    "    images = train_data[idxs]\n",
    "    show_samples(images, title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_HLQR8FikQd_"
   },
   "source": [
    "# Task 2: VAE with Autoregressive flow-based prior on CIFAR10 (5pt)\n",
    "\n",
    "In this task you will fit the VAE model with [flow-based prior](https://arxiv.org/abs/1611.02731) to the CIFAR10 dataset (download it [here](https://drive.google.com/file/d/16j3nrJV821VOkkuRz7aYam8TyIXLnNme/view?usp=sharing)).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 264
    },
    "id": "dTW5lhwv-J9w",
    "outputId": "677167ce-7735-4b9c-ae1b-f61191a0b0cd"
   },
   "outputs": [],
   "source": [
    "train_data, test_data = load_pickle(os.path.join('drive', 'MyDrive', 'DGM', 'homework_supplementary', 'cifar10.pkl'))\n",
    "visualize_data(train_data, 'CIFAR10 samples')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EMNE7poKkqio"
   },
   "source": [
    "The model consists of:\n",
    "* convolutional encoder (variational posterior destrituion $q(\\mathbf{z} | \\mathbf{x})$)\n",
    "* convolutional decoder $p(\\mathbf{x} | \\mathbf{z})$\n",
    "* autoregressive prior\n",
    "\n",
    "We will use MADE model for autoregressive prior. MADE Autoregressive frow (mapping from $\\mathbf{z}\\rightarrow \\boldsymbol{\\epsilon}$) should output location $\\mu(\\mathbf{z})$ and scale parameters $\\sigma(\\mathbf{z})$. The mapping from $\\mathbf{z}$ to $\\boldsymbol{\\epsilon}$ has the form:\n",
    "$$\n",
    "    \\boldsymbol{\\epsilon} = \\mathbf{z} * \\sigma(\\mathbf{z}) + \\mu(\\mathbf{z}).\n",
    "$$\n",
    "\n",
    "The ELBO objective in this task is:\n",
    "$$\n",
    "    -E_{\\mathbf{z}\\sim q(\\mathbf{z}|\\mathbf{x})}[\\log{p(x|\\mathbf{z})}] + E_{z\\sim q(\\mathbf{z}|\\mathbf{x})}[\\log{q(\\mathbf{z}|\\mathbf{x})} - \\log{p(\\mathbf{z})}]\n",
    "$$\n",
    "where \n",
    "$$\n",
    "    \\log{p(\\mathbf{z})} = \\log{p(\\boldsymbol{\\epsilon})} + \\log{\\det\\left|\\frac{d\\boldsymbol{\\epsilon}}{d\\mathbf{z}}\\right|}.\n",
    "$$\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W3AxLZam9cb-"
   },
   "outputs": [],
   "source": [
    "class ConvEncoder(nn.Module):\n",
    "    # do not change this class\n",
    "    def __init__(self, input_shape, n_latent):\n",
    "        super().__init__()\n",
    "        self.input_shape = input_shape\n",
    "        self.n_latent = n_latent\n",
    "        self.convs = nn.Sequential(\n",
    "            nn.Conv2d(input_shape[0], 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 64, 3, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(64, 128, 3, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(128, 256, 3, stride=2, padding=1),\n",
    "        )\n",
    "        conv_out_dim = input_shape[1] // 8 * input_shape[2] // 8 * 256\n",
    "        self.fc = nn.Linear(conv_out_dim, 2 * n_latent)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.convs(x)\n",
    "        out = out.view(out.shape[0], -1)\n",
    "        mu, log_std = self.fc(out).chunk(2, dim=1)\n",
    "        return mu, log_std\n",
    "        \n",
    "\n",
    "class ConvDecoder(nn.Module):\n",
    "    # do not change this class\n",
    "    def __init__(self, n_latent, output_shape):\n",
    "        super().__init__()\n",
    "        self.n_latent = n_latent\n",
    "        self.output_shape = output_shape\n",
    "\n",
    "        self.base_size = (128, output_shape[1] // 8, output_shape[2] // 8)\n",
    "        self.fc = nn.Linear(n_latent, np.prod(self.base_size))\n",
    "        self.deconvs = nn.Sequential(\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(128, 128, 4, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, output_shape[0], 3, padding=1),\n",
    "        )\n",
    "\n",
    "    def forward(self, z):\n",
    "        out = self.fc(z)\n",
    "        out = out.view(out.shape[0], *self.base_size)\n",
    "        return self.deconvs(out)\n",
    "\n",
    "\n",
    "class MaskedLinear(nn.Linear):\n",
    "    # do not change this class\n",
    "    def __init__(self, in_features, out_features, bias=True):\n",
    "        super().__init__(in_features, out_features, bias)\n",
    "        self.register_buffer('mask', torch.ones(out_features, in_features))\n",
    "\n",
    "    def set_mask(self, mask):\n",
    "        self.mask.data.copy_(torch.from_numpy(mask.astype(np.uint8).T))\n",
    "\n",
    "    def forward(self, input):\n",
    "        return F.linear(input, self.mask * self.weight, self.bias)\n",
    "\n",
    "\n",
    "class MADE(nn.Module):\n",
    "    # do not change this class\n",
    "    def __init__(self, input_shape, d, hidden_size=[512, 512]):\n",
    "        super().__init__()\n",
    "        self.input_shape = input_shape\n",
    "        self.nin = np.prod(input_shape)\n",
    "        self.nout = self.nin * d\n",
    "        self.d = d\n",
    "        self.hidden_sizes = hidden_size\n",
    "        self.ordering = np.arange(self.nin)\n",
    "\n",
    "        self.net = []\n",
    "        hs = [self.nin] + self.hidden_sizes + [self.nout]\n",
    "        for h0, h1 in zip(hs, hs[1:]):\n",
    "            self.net.extend([\n",
    "                MaskedLinear(h0, h1),\n",
    "                nn.ReLU(),\n",
    "            ])\n",
    "        self.net.pop()\n",
    "        self.net = nn.ModuleList(self.net)\n",
    "\n",
    "        self.m = {}\n",
    "        self.create_mask()\n",
    "\n",
    "    def create_mask(self):\n",
    "        L = len(self.hidden_sizes)\n",
    "\n",
    "        self.m[-1] = self.ordering\n",
    "        for l in range(L):\n",
    "            self.m[l] = np.random.randint(self.m[l - 1].min(),\n",
    "                                          self.nin - 1, size=self.hidden_sizes[l])\n",
    "\n",
    "        masks = [self.m[l - 1][:, None] <= self.m[l][None, :] for l in range(L)]\n",
    "        masks.append(self.m[L - 1][:, None] < self.m[-1][None, :])\n",
    "\n",
    "        masks[-1] = np.repeat(masks[-1], self.d, axis=1)\n",
    "\n",
    "        layers = [l for l in self.net.modules() if isinstance(l, MaskedLinear)]\n",
    "        for l, m in zip(layers, masks):\n",
    "            l.set_mask(m)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        out = x.view(batch_size, self.nin)\n",
    "        for layer in self.net:\n",
    "            out = layer(out)\n",
    "        out = out.view(batch_size, self.nin, self.d)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ARFPriorVAE(nn.Module):\n",
    "    def __init__(self, input_shape, n_latent):\n",
    "        super().__init__()\n",
    "        assert len(input_shape) == 3\n",
    "        self.input_shape = input_shape\n",
    "        self.n_latent = n_latent\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # define made model, encoder and decoder\n",
    "        # ====\n",
    "\n",
    "    def loss(self, x):\n",
    "        x = 2 * x.float() - 1\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) apply encoder to x to get variational distribution parameters\n",
    "        # 2) sample z from variational distribution\n",
    "        # 3) apply decoder to get reconstruction\n",
    "        \n",
    "        \n",
    "        # ====\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) compute reconstruction loss \n",
    "        # in this case we could use mse loss \n",
    "        # (we will get beta-VAE model since the contributions of reconstruction loss and KL term become dishonest) \n",
    "        # 2) compute encoder log prob (it is a log of normal distribution on z)\n",
    "        # 3) apply MADE model to z to get mu and log_std\n",
    "\n",
    "        # ====\n",
    "\n",
    "        # this trick is just for model stability (do not touch it)\n",
    "        log_std = torch.tanh(log_std)\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) scale z to sigma and shift to mu get epsilon\n",
    "        # 2) compute prior log prob (log of standart normal)\n",
    "        # 3) kl loss is difference between encoder log prob and prior log prob\n",
    "\n",
    "        # ====\n",
    "        return {\n",
    "            'total_loss': recon_loss + kl_loss,\n",
    "            'recon_loss': recon_loss,\n",
    "            'kl_loss': kl_loss\n",
    "        }\n",
    "\n",
    "    def sample(self, n):\n",
    "        with torch.no_grad():\n",
    "            z = torch.randn(n, self.n_latent).cuda()\n",
    "            # investigate how to sample from autoregressive model (do not change this part)\n",
    "            for i in range(self.n_latent):\n",
    "                mu, log_std = self.made(z)[:, i].chunk(2, dim=-1)\n",
    "                log_std = torch.tanh(log_std)\n",
    "                mu, log_std = mu.squeeze(-1), log_std.squeeze(-1)\n",
    "                z[:, i] = (z[:, i] - mu) * torch.exp(-log_std)\n",
    "            return self.decoder(z).cpu().numpy() * 0.5 + 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 311,
     "referenced_widgets": [
      "ceaf984c21cb450fa03174b736da865f",
      "7472906555e342dc81c967fdac6813b4",
      "bad84819d77244559184f00b46493fd6",
      "4eec5b3b5d7947329e2c55a01c90fe72",
      "b83f8b841c02424dabc60cc3d8980bde",
      "e58c8556e1cc4f56838a0683679135cf",
      "15d0dc877bda48a29da18aa1efbca800",
      "a70d374e86314adeae856a9dfe5519f3",
      "d7f9660b4f234c9380aa68e165c73e06",
      "e01506f4adb749b7af1b95b439d4b108",
      "2ad6eb2c745545dd912ce97ee9110f11"
     ]
    },
    "id": "JcMArVCA9ckD",
    "outputId": "1dd59fc5-11f2-433c-e6c7-2d969bcdff33"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters\n",
    "BATCH_SIZE = \n",
    "EPOCHS = \n",
    "LR = \n",
    "# ====\n",
    "\n",
    "train_data, test_data = load_pickle(os.path.join('drive', 'MyDrive', 'DGM', 'homework_supplementary', 'cifar10.pkl'))\n",
    "model = ARFPriorVAE((3, 32, 32), 16)\n",
    "\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = data.DataLoader(test_data, batch_size=BATCH_SIZE)\n",
    "train_losses, test_losses = train_model(model, train_loader, test_loader, use_cuda=USE_CUDA, epochs=EPOCHS, use_tqdm=True, lr=LR)\n",
    "\n",
    "plot_training_curves(train_losses, test_losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 812
    },
    "id": "beJCpjz0edsW",
    "outputId": "ebe8ad16-5e0e-4669-df01-a1da795a486f"
   },
   "outputs": [],
   "source": [
    "samples = model.sample(100)\n",
    "\n",
    "x = next(iter(test_loader))[:50].cuda()\n",
    "with torch.no_grad():\n",
    "    z = model.encoder(2 * x - 1)[0]\n",
    "    x_recon = model.decoder(z).cpu().numpy() * 0.5 + 0.5\n",
    "x = x.cpu().numpy()\n",
    "reconstructions = np.concatenate((x, x_recon), axis=0)\n",
    "\n",
    "x = next(iter(test_loader))[:20].cuda()\n",
    "with torch.no_grad():\n",
    "    x = 2 * x - 1\n",
    "    z, _ = model.encoder(x)\n",
    "    z1, z2 = z.chunk(2, dim=0)\n",
    "    interps = [model.decoder(z1 * (1 - alpha) + z2 * alpha) for alpha in np.linspace(0, 1, 10)]\n",
    "    interps = torch.stack(interps, dim=1).view(-1, 3, 32, 32)\n",
    "    interps = torch.clamp(interps, -1, 1) * 0.5 + 0.5\n",
    "interps = interps.cpu().numpy()\n",
    "\n",
    "samples = np.clip(samples, 0, 1)\n",
    "reconstructions = np.clip(reconstructions, 0, 1)\n",
    "interps = np.clip(interps, 0, 1)\n",
    "\n",
    "for key, value in test_losses.items():\n",
    "    print('{}: {:.4f}'.format(key, value[-1]))\n",
    "    \n",
    "show_samples(samples, title='Samples')\n",
    "show_samples(reconstructions, title='Reconstructions')\n",
    "show_samples(interps, title='Interpolations')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qZZ-ef7XkQiT"
   },
   "source": [
    "# Task 3: VAE with Autoregressive decoder on MNIST (5pt)\n",
    "\n",
    "In this task you will fit the VAE model with [autoregressive decoder](https://arxiv.org/abs/1611.05013) to the MNIST dataset (download it [here](https://drive.google.com/file/d/1Ms-RBybrueI3_w2CRj7lM9mYjfvFRL6w/view?usp=sharing))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 264
    },
    "id": "d7SDW4JHgKFr",
    "outputId": "fe7cfeaf-18cf-4258-97f6-84d078383c95"
   },
   "outputs": [],
   "source": [
    "# change the path to the file\n",
    "train_data, test_data = load_pickle(os.path.join('drive', 'My Drive', 'DGM', 'homework_supplementary', 'mnist.pkl'), binarize=True)\n",
    "visualize_data(train_data, 'MNIST samples')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iazovzdJpc2Q"
   },
   "outputs": [],
   "source": [
    "class MaskedConv2d(nn.Conv2d):\n",
    "    def __init__(self, mask_type, in_channels, out_channels, kernel_size=5, padding=0, conditional_size=None):\n",
    "        assert mask_type in ['A', 'B']\n",
    "        super().__init__(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=padding)\n",
    "        self.register_buffer('mask', torch.zeros_like(self.weight))\n",
    "        self.create_mask(mask_type)\n",
    "\n",
    "        if conditional_size is not None:\n",
    "            self.cond_op = nn.Linear(conditional_size, self.out_channels)\n",
    "\n",
    "    def forward(self, input, cond=None):\n",
    "        # ====\n",
    "        # your code\n",
    "        # apply masked convolution and get \"out\" variable\n",
    "\n",
    "        # ====\n",
    "\n",
    "        if cond is not None:\n",
    "            cond = self.cond_op(cond)\n",
    "            out = out + cond.view(cond.shape[0], self.out_channels, 1, 1)\n",
    "        return out\n",
    "\n",
    "    def create_mask(self, mask_type):\n",
    "        # ====\n",
    "        # your code\n",
    "        # do not forget about mask_type\n",
    "        \n",
    "        # ====\n",
    "\n",
    "\n",
    "def test_masked_conv2d():\n",
    "    layer = MaskedConv2d('A', 2, 2)\n",
    "    assert np.allclose(layer.mask[:, :, 2, 2].numpy(), np.zeros((2, 2)))\n",
    "\n",
    "    layer = MaskedConv2d('B', 2, 2)\n",
    "    assert np.allclose(layer.mask[:, :, 2, 2].numpy(), np.ones((2, 2)))\n",
    "\n",
    "\n",
    "test_masked_conv2d()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Pb0FcBizpYLO"
   },
   "source": [
    "[Layer Normalization](https://arxiv.org/abs/1607.06450) helps to stabilize training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o0IbeNTGpWIx"
   },
   "outputs": [],
   "source": [
    "class LayerNorm(nn.LayerNorm):\n",
    "    def __init__(self, n_filters):\n",
    "        super().__init__(n_filters)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.permute(0, 2, 3, 1).contiguous()\n",
    "        x = super().forward(x)\n",
    "        return x.permute(0, 3, 1, 2).contiguous()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Sg953_qZOmGR"
   },
   "outputs": [],
   "source": [
    "def get_normal_KL(mean_1, log_std_1, mean_2=None, log_std_2=None):\n",
    "    \"\"\"\n",
    "        This function should return the value of KL(p1 || p2),\n",
    "        where p1 = Normal(mean_1, exp(log_std_1)), p2 = Normal(mean_2, exp(log_std_2)).\n",
    "        If mean_2 and log_std_2 are None values, we will use standart normal distribution.\n",
    "        Note that we consider the case of diagonal covariance matrix.\n",
    "    \"\"\"\n",
    "    if mean_2 is None:\n",
    "        mean_2 = torch.zeros_like(mean_1)\n",
    "    if log_std_2 is None:\n",
    "        log_std_2 = torch.zeros_like(log_std_1)\n",
    "    # ====\n",
    "    # your code\n",
    "    \n",
    "    # ====\n",
    "\n",
    "\n",
    "def test_KL():\n",
    "    assert np.isclose(get_normal_KL(torch.tensor(2), torch.tensor(3), torch.tensor(0), torch.tensor(0)).numpy(), 200.2144, rtol=1e-3)\n",
    "    assert np.isclose(get_normal_KL(torch.tensor(2), torch.tensor(3), torch.tensor(4), torch.tensor(5)).numpy(), 1.50925, rtol=1e-3)\n",
    "    assert np.allclose(get_normal_KL(torch.tensor((10, 10)), torch.tensor((2, 4)), torch.tensor((3, 5))).numpy(), [49.2990, 1498.479], rtol=1e-3)\n",
    "\n",
    "test_KL()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I4J1G9LDCerC"
   },
   "outputs": [],
   "source": [
    "class PixelCNN(nn.Module):\n",
    "    def __init__(\n",
    "        self, \n",
    "        input_shape, \n",
    "        n_filters=256, \n",
    "        kernel_size=3, \n",
    "        n_layers=7,\n",
    "        use_layer_norm=True,\n",
    "        conditional_size=None\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.input_shape = input_shape\n",
    "        \n",
    "        # ====\n",
    "        # your code\n",
    "        # apply the sequence of MaskedConv2d -> LayerNorm -> ReLU\n",
    "        # note that the first conv layer should be of type 'A'\n",
    "        # the last layer should be MaskedConv2d\n",
    "        # define self.net as list of layers\n",
    "        \n",
    "\n",
    "        # ====\n",
    "\n",
    "    def forward(self, x, cond=None):\n",
    "        out = (x.float() - 0.5) * 2\n",
    "        for layer in self.net:\n",
    "            if isinstance(layer, MaskedConv2d):\n",
    "                out = layer(out, cond=cond)\n",
    "            else:\n",
    "                out = layer(out)\n",
    "        return out.view(x.shape[0], 2, 1, *self.input_shape)\n",
    "\n",
    "    def loss(self, x, cond=None):\n",
    "        # ====\n",
    "        # your code\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def sample(self, n, cond=None):\n",
    "        # read carefully the sampling process\n",
    "        samples = torch.zeros(n, 1, *self.input_shape).cuda()\n",
    "        with torch.no_grad():\n",
    "            for r in range(self.input_shape[0]):\n",
    "                for c in range(self.input_shape[1]):\n",
    "                    logits = self(samples)[:, :, :, r, c]\n",
    "                    probs = F.softmax(logits, dim=1).squeeze(-1)\n",
    "                    samples[:, 0, r, c] = torch.multinomial(probs, num_samples=1).squeeze(-1)\n",
    "        return samples.cpu().numpy()\n",
    "\n",
    "\n",
    "class ConvEncoder(nn.Module):\n",
    "    def __init__(self, input_shape, latent_dim):\n",
    "        super().__init__()\n",
    "        self.input_shape = input_shape\n",
    "        self.latent_dim = latent_dim\n",
    "        self.convs = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 32, 3, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 64, 3, stride=2, padding=1),\n",
    "        )\n",
    "        conv_out_dim = input_shape[0] // 4 * input_shape[1] // 4 * 64\n",
    "        self.fc = nn.Linear(conv_out_dim, 2 * latent_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = (x.float() - 0.5) * 2\n",
    "        out = self.convs(x)\n",
    "        out = out.view(out.shape[0], -1)\n",
    "        mu, log_std = self.fc(out).chunk(2, dim=1)\n",
    "        return mu, log_std\n",
    "\n",
    "\n",
    "class ARDecoderVAE(nn.Module):\n",
    "    def __init__(self, input_shape, n_latent, free_bits=None):\n",
    "        super().__init__()\n",
    "        assert len(input_shape) == 2\n",
    "\n",
    "        self.input_shape = input_shape\n",
    "        self.n_latent = n_latent\n",
    "        self.free_bits = free_bits\n",
    "        self.encoder = ConvEncoder(input_shape, n_latent)\n",
    "        self.decoder = PixelCNN(\n",
    "            input_shape, \n",
    "            n_filters=32, \n",
    "            n_layers=3,\n",
    "            kernel_size=5, \n",
    "            conditional_size=n_latent\n",
    "        )\n",
    "\n",
    "    def prior(self, n):\n",
    "        # ====\n",
    "        # your code\n",
    "        # return n samples from prior distribution (we use standart normal for prior)\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def loss(self, x):\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) apply encoder\n",
    "        # 2) apply reparametrization trick\n",
    "        # 3) get decoder loss (reconstruction loss)\n",
    "        # 4) get kl loss using get_normal_KL\n",
    "        # 5) apply free_bits\n",
    "        \n",
    "        # ====\n",
    "\n",
    "        return {\n",
    "            'total_loss': recon_loss + kl_loss, \n",
    "            'recon_loss': recon_loss,\n",
    "            'kl_loss': kl_loss\n",
    "        }\n",
    "\n",
    "    def sample(self, n):\n",
    "        with torch.no_grad():\n",
    "            z = self.prior(n).cuda()\n",
    "            samples = self.decoder.sample(n, cond=z)\n",
    "        return samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "8ccdc2d1ce684322bf1cd66703db0ec1",
      "3d9a2dee58b94433992c58a3a4f3928a",
      "69a166fbb1a0405ab770967d79281431",
      "9dff0d3e3ec84df69e9e9f194a174c43",
      "34ad9afa593e47b7b43a6434bcaf67c8",
      "21e4ffb3a1c54713a859362b4859e623",
      "94d8ea7e94db469bb7a562801b3ea55c",
      "ce058edc7d434b1a9f680d242e9f8faa",
      "bd656738a3524b7e80261dbd04c19643",
      "61bf6ab3645b4b97a1b4210acd909e29",
      "6238d43942914e498da5faaaf1c5a848"
     ]
    },
    "id": "i4IgkdDuCetl",
    "outputId": "9d0a1ee3-d82a-4efb-fd1c-492d66a1c1a2"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters\n",
    "BATCH_SIZE = \n",
    "EPOCHS = \n",
    "LR = \n",
    "FREE_BITS = \n",
    "# ====\n",
    "\n",
    "train_data, test_data = load_pickle(os.path.join('drive', 'My Drive', 'DGM', 'homework_supplementary', 'mnist.pkl'), binarize=True)\n",
    "\n",
    "model = ARDecoderVAE(input_shape=(28, 28), free_bits=FREE_BITS, n_latent=16)\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = data.DataLoader(test_data, batch_size=BATCH_SIZE)\n",
    "\n",
    "train_losses, test_losses = train_model(model, train_loader, test_loader, epochs=EPOCHS, use_cuda=USE_CUDA, use_tqdm=True, lr=LR)\n",
    "\n",
    "plot_training_curves(train_losses, test_losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 827
    },
    "id": "E_X3IGjWCew0",
    "outputId": "a63b447a-208e-44fa-a4a0-d0e5f46e9692"
   },
   "outputs": [],
   "source": [
    "samples = model.sample(100)\n",
    "\n",
    "x = next(iter(test_loader))[:50].cuda()\n",
    "with torch.no_grad():\n",
    "    z, _ = model.encoder(2 * x - 1)\n",
    "    x_recon = model.decoder.sample(50, cond=z)\n",
    "x = x.cpu().numpy()\n",
    "reconstructions = np.concatenate((x, x_recon), axis=0)\n",
    "\n",
    "samples = samples.astype('float32')\n",
    "reconstructions = reconstructions.astype('float32')\n",
    "\n",
    "for key, value in test_losses.items():\n",
    "    print('{}: {:.4f}'.format(key, value[-1]))\n",
    "show_samples(samples, title='Samples')\n",
    "show_samples(reconstructions, title='Reconstructions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ijciVfzeno1M"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "hw4.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "15d0dc877bda48a29da18aa1efbca800": {
     "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
     }
    },
    "21e4ffb3a1c54713a859362b4859e623": {
     "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": ""
     }
    },
    "2ad6eb2c745545dd912ce97ee9110f11": {
     "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
     }
    },
    "34ad9afa593e47b7b43a6434bcaf67c8": {
     "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_6238d43942914e498da5faaaf1c5a848",
      "placeholder": "​",
      "style": "IPY_MODEL_61bf6ab3645b4b97a1b4210acd909e29",
      "value": " 3/3 [03:25&lt;00:00, 68.46s/it]"
     }
    },
    "3d9a2dee58b94433992c58a3a4f3928a": {
     "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
     }
    },
    "4eec5b3b5d7947329e2c55a01c90fe72": {
     "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_d7f9660b4f234c9380aa68e165c73e06",
      "max": 3,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_a70d374e86314adeae856a9dfe5519f3",
      "value": 3
     }
    },
    "61bf6ab3645b4b97a1b4210acd909e29": {
     "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": ""
     }
    },
    "6238d43942914e498da5faaaf1c5a848": {
     "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
     }
    },
    "69a166fbb1a0405ab770967d79281431": {
     "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_94d8ea7e94db469bb7a562801b3ea55c",
      "placeholder": "​",
      "style": "IPY_MODEL_21e4ffb3a1c54713a859362b4859e623",
      "value": "100%"
     }
    },
    "7472906555e342dc81c967fdac6813b4": {
     "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
     }
    },
    "8ccdc2d1ce684322bf1cd66703db0ec1": {
     "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_69a166fbb1a0405ab770967d79281431",
       "IPY_MODEL_9dff0d3e3ec84df69e9e9f194a174c43",
       "IPY_MODEL_34ad9afa593e47b7b43a6434bcaf67c8"
      ],
      "layout": "IPY_MODEL_3d9a2dee58b94433992c58a3a4f3928a"
     }
    },
    "94d8ea7e94db469bb7a562801b3ea55c": {
     "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
     }
    },
    "9dff0d3e3ec84df69e9e9f194a174c43": {
     "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_bd656738a3524b7e80261dbd04c19643",
      "max": 3,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_ce058edc7d434b1a9f680d242e9f8faa",
      "value": 3
     }
    },
    "a70d374e86314adeae856a9dfe5519f3": {
     "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": ""
     }
    },
    "b83f8b841c02424dabc60cc3d8980bde": {
     "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_2ad6eb2c745545dd912ce97ee9110f11",
      "placeholder": "​",
      "style": "IPY_MODEL_e01506f4adb749b7af1b95b439d4b108",
      "value": " 3/3 [01:49&lt;00:00, 36.17s/it]"
     }
    },
    "bad84819d77244559184f00b46493fd6": {
     "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_15d0dc877bda48a29da18aa1efbca800",
      "placeholder": "​",
      "style": "IPY_MODEL_e58c8556e1cc4f56838a0683679135cf",
      "value": "100%"
     }
    },
    "bd656738a3524b7e80261dbd04c19643": {
     "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
     }
    },
    "ce058edc7d434b1a9f680d242e9f8faa": {
     "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": ""
     }
    },
    "ceaf984c21cb450fa03174b736da865f": {
     "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_bad84819d77244559184f00b46493fd6",
       "IPY_MODEL_4eec5b3b5d7947329e2c55a01c90fe72",
       "IPY_MODEL_b83f8b841c02424dabc60cc3d8980bde"
      ],
      "layout": "IPY_MODEL_7472906555e342dc81c967fdac6813b4"
     }
    },
    "d7f9660b4f234c9380aa68e165c73e06": {
     "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
     }
    },
    "e01506f4adb749b7af1b95b439d4b108": {
     "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": ""
     }
    },
    "e58c8556e1cc4f56838a0683679135cf": {
     "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": ""
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}