homeworks/hw5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mg5oRBoOWloW"
   },
   "source": [
    "## Task 1: Theory (3pt)\n",
    "\n",
    "1. Let's recall the variational lower bound for the discrete data dequantization problem from Lecture 7:\n",
    "$$\n",
    "    \\log P(\\mathbf{x} | \\boldsymbol{\\theta}) \\geq  \\int q(\\mathbf{u} | \\mathbf{x}) \\log \\frac{p(\\mathbf{x} + \\mathbf{u} | \\boldsymbol{\\theta})}{q(\\mathbf{u} | \\mathbf{x})} d \\mathbf{u} = \\mathcal{L}(q, \\boldsymbol{\\theta}).\n",
    "$$\n",
    "We have discussed with you that the variational lower bound can be improved with the help of Importance Sampling technique. Write out the lower bound for dequantization using $\\mathcal{L}_k$ by analogy with the IWAE model (you need to use not just one $\\mathbf{u}$, but a set of $\\{\\mathbf{u}_k\\}_{k=1}^K$).\n",
    "    \n",
    "2. The Vanilla GAN often suffers from problems with a vanishing gradient. [Least Squares GAN](https://arxiv.org/abs/1611.04076) tries to solve this problem by replacing the error function with the following:\n",
    "$$\n",
    "   \t\\min_D V(D) = \\min_D \\frac{1}{2}\\left[ \\mathbb{E}_{\\pi(\\mathbf{x})} (D(\\mathbf{x}) - b)^2 + \\mathbb{E}_{p(\\mathbf{z})} (D(G(\\mathbf{z})) - a)^2 \\right]\n",
    "$$\n",
    "$$\n",
    "   \t\\min_G V(G) = \\min_G \\frac{1}{2}\\left[ \\mathbb{E}_{\\pi(\\mathbf{x})} (D(\\mathbf{x}) - c)^2 + \\mathbb{E}_{p(\\mathbf{z})} (D(G(\\mathbf{z})) - c)^2 \\right],\n",
    "$$\n",
    "where $a,b,c \\in \\mathbb{R}$ some fixed constants.\n",
    "    * Find the formula for the optimal discriminator $D^*$.\n",
    "    * Write out the expression for the error function of the generator $V(G)$ in the case of an optimal discriminator $D^*$.\n",
    "    * Prove that for $b - c = 1$, $b - a = 2$, the error function of the generator $V(G)$ in the case of the optimal discriminator $D^*$ takes the form:\n",
    "$$\n",
    "   \tV(G) = \\chi^2_{\\text{Pearson}} \\left(\\frac{\\pi(\\mathbf{x}) + p(\\mathbf{x} | \\boldsymbol{\\theta})}{2} || p(\\mathbf{x} | \\boldsymbol{\\theta})\\right), \n",
    "$$\n",
    "where $\\chi^2_{\\text{Pearson}} (p || q)$ is a squared Pearson divergence:\n",
    "$$\n",
    "   \t\\chi^2_{\\text{Pearson}} (p || q) = \\int \\frac{(p(\\mathbf{x}) - q(\\mathbf{x}))^2}{p(\\mathbf{x})} d \\mathbf{x}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kwS2oTkaVLct"
   },
   "source": [
    "```your solution```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "fJQYtKgA5ate",
    "outputId": "dc4083a6-2bea-4e10-da5e-20d1f866f352"
   },
   "outputs": [],
   "source": [
    "REPO_NAME = \"2022-DGM-Ozon-course\"\n",
    "!if [ -d {REPO_NAME} ]; then rm -Rf {REPO_NAME}; fi\n",
    "!git clone https://github.com/r-isachenko/{REPO_NAME}.git\n",
    "!cd {REPO_NAME}\n",
    "!pip install ./{REPO_NAME}/homeworks/\n",
    "!rm -Rf {REPO_NAME}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Sl59dviN5ekr"
   },
   "outputs": [],
   "source": [
    "from dgm_utils import visualize_2d_samples, show_samples, visualize_images, load_pickle "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qCNBV60b8CHX",
    "outputId": "9c18547f-fd3c-4a8f-bacc-776fdae68e6e"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "\n",
    "print('cuda is available:', USE_CUDA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mQb5MtUKd_ii",
    "outputId": "6a9aafac-1956-45b0-cf07-56c4ceeea1f7"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ftMgDK_cPFmD"
   },
   "outputs": [],
   "source": [
    "# do not change this function\n",
    "def plot_losses(losses, title):\n",
    "    n_itr = len(losses)\n",
    "    xs = np.arange(n_itr)\n",
    "\n",
    "    plt.figure(figsize=(7, 5))\n",
    "    plt.plot(xs, losses)\n",
    "    plt.title(title, fontsize=14)\n",
    "    plt.xlabel('Iterations', fontsize=14)\n",
    "    plt.ylabel('Loss', fontsize=14)\n",
    "\n",
    "    plt.xticks(fontsize=12)\n",
    "    plt.yticks(fontsize=12)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eF2vatv9ikAi"
   },
   "source": [
    "# Task 2: WGAN vs WGAN-GP on 2d data (5pt)\n",
    "\n",
    "Here your task is to reproduce the experiment with a toy dataset from [WGAN-GP](https://arxiv.org/pdf/1704.00028.pdf) paper.\n",
    "\n",
    "Let generate data. Our data is a set of gaussians. Read the following snippet of code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VTkzYqkufT-W"
   },
   "outputs": [],
   "source": [
    "def generate_2d_data(size):\n",
    "    scale = 2\n",
    "    var = 0.02\n",
    "    centers = [\n",
    "        (1, 0),\n",
    "        (-1, 0),\n",
    "        (0, 1),\n",
    "        (0, -1),\n",
    "        (1. / np.sqrt(2), 1. / np.sqrt(2)),\n",
    "        (1. / np.sqrt(2), -1. / np.sqrt(2)),\n",
    "        (-1. / np.sqrt(2), 1. / np.sqrt(2)),\n",
    "        (-1. / np.sqrt(2), -1. / np.sqrt(2))\n",
    "    ]\n",
    "    \n",
    "    centers = [(scale * x, scale * y) for x, y in centers]\n",
    "    dataset = []\n",
    "    \n",
    "    for i in range(size):    \n",
    "        point = np.random.randn(2) * var\n",
    "        center = centers[np.random.choice(np.arange(len(centers)))]\n",
    "        point[0] += center[0]\n",
    "        point[1] += center[1]\n",
    "        dataset.append(point)\n",
    "        \n",
    "    dataset = np.array(dataset, dtype='float32')\n",
    "    dataset /= 1.414  # stdev\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NSa9-5sQeme6"
   },
   "source": [
    "The data has lots of separate modes. \n",
    "\n",
    "Our goal is to compare WGAN and WGAN-GP models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 341
    },
    "id": "pB58-3K7iuXs",
    "outputId": "664dbb92-0ceb-4054-869e-76d68bb69e1f"
   },
   "outputs": [],
   "source": [
    "COUNT = 20000\n",
    "\n",
    "train_data = generate_2d_data(COUNT)\n",
    "visualize_2d_samples(train_data, \"Train data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H7D48JfdjRCG"
   },
   "source": [
    "We will use the following helper functions. Please, read them carefully!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lN8eEktgipYY"
   },
   "outputs": [],
   "source": [
    "def make_inference(generator, critic, n_samples=5000):\n",
    "    generator.eval()\n",
    "    critic.eval()\n",
    "    xs = np.linspace(-3.0, 3.0, 1000 + 1)\n",
    "    xg, yg = np.meshgrid(xs, xs)\n",
    "    grid = np.concatenate((xg.reshape(-1, 1), yg.reshape(-1, 1)), axis=-1)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        samples = generator.sample(n_samples).cpu().detach().numpy()\n",
    "        critic_output = critic(torch.FloatTensor(grid).cuda()).cpu().detach().numpy()\n",
    "    \n",
    "    critic_output = np.prod(critic_output, axis=-1).reshape((1000 + 1, 1000 + 1))\n",
    "    return samples, grid, critic_output\n",
    "\n",
    "\n",
    "def visualize_critic_output(generated_samples, real_samples, grid, critic_output, npts=100 + 1):\n",
    "    plt.figure(figsize=(6, 6))\n",
    "    plt.gca().set_aspect(\"equal\")\n",
    "    \n",
    "    npts = critic_output.shape[0]\n",
    "    cnt = plt.contourf(\n",
    "        grid[:, 0].reshape((npts, npts)), grid[:, 1].reshape((npts, npts)), critic_output,\n",
    "        levels=25, cmap=\"cividis\"\n",
    "    )\n",
    "    plt.scatter(generated_samples[:, 0], generated_samples[:, 1], marker=\".\", color=\"red\", s=0.5)\n",
    "    plt.scatter(real_samples[:, 0], real_samples[:, 1], marker=\"x\", color=\"blue\", s=0.5)\n",
    "    plt.colorbar(cnt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8klUHHmAjHsI"
   },
   "source": [
    "## WGAN\n",
    "\n",
    "[WGAN](https://arxiv.org/abs/1701.07875) model uses weight clipping to enforce Lipschitzness of the critic.\n",
    "\n",
    "The model objective is\n",
    "$$\n",
    "\\min_{G} W(\\pi || p) \\approx \\min_{G} \\max_{\\boldsymbol{\\phi} \\in \\boldsymbol{\\Phi}} \\left[ \\mathbb{E}_{\\pi(\\mathbf{x})} f(\\mathbf{x}, \\boldsymbol{\\phi})  - \\mathbb{E}_{p(\\mathbf{z})} f(G(\\mathbf{z}, \\boldsymbol{\\theta}), \\boldsymbol{\\phi} )\\right].\n",
    "$$\n",
    "Here $f(\\mathbf{x}, \\boldsymbol{\\phi})$ is the critic model. The critic weights $\\boldsymbol{\\phi}$ should lie in the compact set $\\boldsymbol{\\Phi} = [-c, c]^d$.\n",
    "\n",
    "In this task we will use fully-connected networks for the generator $G(\\mathbf{z}, \\boldsymbol{\\theta})$ and the critic $f(\\mathbf{x}, \\boldsymbol{\\phi})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "daZ0rPpXePCX"
   },
   "outputs": [],
   "source": [
    "class FullyConnectedMLP(nn.Module):\n",
    "    def __init__(self, input_dim, hiddens, output_dim):\n",
    "        assert isinstance(hiddens, list)\n",
    "        super().__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.hiddens = hiddens\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # stack Linear layear and ReLU (the last layer shoud be Linear)\n",
    "        self.net = \n",
    "        # ====\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        x = x.view(batch_size, -1)\n",
    "        return self.net(x).view(batch_size, self.output_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5S0ihfmv0XI0"
   },
   "source": [
    "Let define the generator and critic networks. \n",
    "\n",
    "We only need to define \n",
    "1. the sample method for the generator;\n",
    "2. weight clipping for the critic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "j1JA7tyweLX9"
   },
   "outputs": [],
   "source": [
    "class MLPGenerator(FullyConnectedMLP):\n",
    "    def sample(self, n):\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) sample from standard normal distribution \n",
    "        # 2) apply the model\n",
    "        \n",
    "        # ===="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fHNt8_5pmdYl"
   },
   "outputs": [],
   "source": [
    "class MLPCritic(FullyConnectedMLP):\n",
    "\n",
    "    def clip_weights(self, c):\n",
    "        for layer in self.net:\n",
    "            if isinstance(layer, nn.Linear):\n",
    "                # ====\n",
    "                # your code\n",
    "                # clip the weight to the range [-c, c]\n",
    "                weight = \n",
    "                # ====\n",
    "                layer.weight = nn.Parameter(weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZAsRaQCD2GeY"
   },
   "source": [
    "Now we need to implement the train loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "echIwLvki0LF"
   },
   "outputs": [],
   "source": [
    "def train_wgan(\n",
    "    generator, \n",
    "    critic, \n",
    "    train_loader,\n",
    "    critic_steps, \n",
    "    batch_size,\n",
    "    n_epochs,\n",
    "    lr, \n",
    "    clip_c,\n",
    "    use_cuda\n",
    "):\n",
    "    train_losses = []\n",
    "\n",
    "    if use_cuda:\n",
    "        generator = generator.cuda()\n",
    "        critic = critic.cuda()\n",
    "\n",
    "    gen_optimizer = torch.optim.Adam(generator.parameters(), lr=lr, betas=(0, 0.9))\n",
    "    critic_optimizer = torch.optim.Adam(critic.parameters(), lr=lr, betas=(0, 0.9))\n",
    "\n",
    "    generator.train()\n",
    "    critic.train()\n",
    "\n",
    "    curr_iter = 0\n",
    "    d_loss, g_loss = torch.zeros(1), torch.zeros(1)\n",
    "    \n",
    "    batch_loss_history = {'discriminator_losses': [], 'generator_losses': []}\n",
    "    for epoch_i in tqdm(range(n_epochs)):\n",
    "        for (batch_i, x) in enumerate(train_loader):\n",
    "            curr_iter += 1\n",
    "            if use_cuda:\n",
    "                x = x.cuda()\n",
    "\n",
    "            # do a critic update\n",
    "            with torch.no_grad():\n",
    "                fake_data = generator.sample(x.shape[0])\n",
    "            \n",
    "            critic_optimizer.zero_grad()\n",
    "            \n",
    "            # ====\n",
    "            # your code\n",
    "            # D(x_fake) - D(x_real)\n",
    "            d_loss = \n",
    "            # ====\n",
    "            d_loss.backward()\n",
    "            critic_optimizer.step()\n",
    "            critic.clip_weights(clip_c)\n",
    "\n",
    "            # generator update\n",
    "            if curr_iter % critic_steps == 0:\n",
    "                gen_optimizer.zero_grad()\n",
    "                fake_data = generator.sample(batch_size)\n",
    "                \n",
    "                # ====\n",
    "                # your code\n",
    "                # -D(x_fake)\n",
    "                g_loss = \n",
    "                # ====\n",
    "                g_loss.backward()\n",
    "                gen_optimizer.step()\n",
    "\n",
    "                batch_loss_history['generator_losses'].append(g_loss.data.cpu().numpy())\n",
    "                batch_loss_history['discriminator_losses'].append(d_loss.data.cpu().numpy())\n",
    "\n",
    "    return batch_loss_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 731,
     "referenced_widgets": [
      "effa9aab4edb4cc18261259a60f7d4e3",
      "68765de1436245f192703b2d139ee0ac",
      "ac937e59ba7f456ab73b2fcf1fe54014",
      "9a51463240d84c21b23600b64cf1a096",
      "bad267b7949340e1860b6475f81ab0de",
      "5b1cb14e8efa4cf5b7be757462f51d74",
      "8345ed2254f244f5bfdbf7bd5bf8c267",
      "84aa41c85adb4e11bfd484e3f316a839",
      "41585a0cbe524f4cbd03022e17393fe6",
      "14179864decd4f2390c42259e9bc426b",
      "d72569fa2f5641fc8ca12367321944fc"
     ]
    },
    "id": "NGesm5J9jKyr",
    "outputId": "b1f4979f-5564-42a4-c6e1-24d57d8e1e11"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters\n",
    "BATCH_SIZE =       # any adequate value\n",
    "GEN_HIDDENS =      # 4 layers with < 128 neurons would be enough\n",
    "DISCR_HIDDENS =    # 4 layers with < 128 neurons would be enough\n",
    "CRITIC_STEPS =     # > 2\n",
    "LR =               # < 1e-2\n",
    "CLIP_C =           # < 1\n",
    "# ====\n",
    "\n",
    "N_EPOCHS = 200 # change it if you want\n",
    "\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "\n",
    "generator = MLPGenerator(input_dim=16, hiddens=GEN_HIDDENS, output_dim=2)\n",
    "critic = MLPCritic(input_dim=2, hiddens=DISCR_HIDDENS, output_dim=1)\n",
    "\n",
    "train_losses = train_wgan(\n",
    "    generator, \n",
    "    critic, \n",
    "    train_loader,\n",
    "    critic_steps=CRITIC_STEPS, \n",
    "    batch_size=BATCH_SIZE, \n",
    "    n_epochs=N_EPOCHS,\n",
    "    lr=LR,\n",
    "    clip_c=CLIP_C,\n",
    "    use_cuda=USE_CUDA\n",
    ")\n",
    "\n",
    "plot_losses(train_losses['discriminator_losses'], 'Critic loss')\n",
    "plot_losses(train_losses['generator_losses'], 'Generator loss')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mPwnnDIR2ZyL"
   },
   "source": [
    "Let plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 363
    },
    "id": "ay6-MjQGjO1V",
    "outputId": "60ba8c5d-0404-486e-bdcf-0d5de103c8aa"
   },
   "outputs": [],
   "source": [
    "samples, grid, critic_output = make_inference(generator, critic)\n",
    "visualize_critic_output(samples, train_data, grid, critic_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZG4ErQLNjZr9"
   },
   "source": [
    "## WGAN-GP\n",
    "\n",
    "[WGAN-GP](https://arxiv.org/pdf/1704.00028.pdf)  model uses gradient penalty to enforce Lipschitzness.\n",
    "\n",
    "The model objective is\n",
    "$$\n",
    "    W(\\pi || p) = \\underbrace{\\mathbb{E}_{\\pi(\\mathbf{x})} f(\\mathbf{x})  - \\mathbb{E}_{p(\\mathbf{x} | \\boldsymbol{\\theta})} f(\\mathbf{x})}_{\\text{original critic loss}} + \\lambda \\underbrace{\\mathbb{E}_{U[0, 1]} \\left[ \\left( \\| \\nabla_{\\hat{\\mathbf{x}}_t} f(\\hat{\\mathbf{x}}_t) \\|_2 - 1 \\right) ^ 2\\right]}_{\\text{gradient penalty}},\n",
    "$$\n",
    "where the samples $\\hat{\\mathbf{x}}_t = t \\mathbf{x} + (1 - t) \\mathbf{y}$ with $t \\in [0, 1]$ are uniformly sampled along straight lines between pairs of points: $\\mathbf{x}$ from the data distribution $\\pi(\\mathbf{x})$ and $\\mathbf{y}$ from the generator distribution $p(\\mathbf{x} | \\boldsymbol{\\theta}))$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qjN5thID3WZ2"
   },
   "source": [
    "Let define our gradient penalty loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mWGSIJ7n3VQk"
   },
   "outputs": [],
   "source": [
    "def gradient_penalty(critic, real_data, fake_data):\n",
    "    batch_size = real_data.shape[0]\n",
    "\n",
    "    # ====\n",
    "    # your code\n",
    "    # Calculate interpolation x_t = t * x_real + (1 - t) x_fake\n",
    "    # 1) sample t\n",
    "    # 2) create x_t (be careful about shapes)\n",
    "    # 3) apply critic to x_t\n",
    "    x_t = \n",
    "    d_output =\n",
    "\n",
    "    # ====\n",
    "\n",
    "    gradients = torch.autograd.grad(\n",
    "        outputs=d_output, \n",
    "        inputs=x_t, \n",
    "        grad_outputs=torch.ones(d_output.size()).to(fake_data.device), \n",
    "        create_graph=True, \n",
    "        retain_graph=True\n",
    "    )[0]\n",
    "\n",
    "    gradients = gradients.reshape(batch_size, -1)\n",
    "    # ====\n",
    "    # your code\n",
    "    # compute gradient norm\n",
    "    gradients_norm = \n",
    "    # ====\n",
    "    return ((gradients_norm - 1) ** 2).mean()\n",
    "\n",
    "\n",
    "def test_gradient_penalty():\n",
    "    x = np.random.normal(size=(10, 4))\n",
    "    x_norm = np.mean(np.sqrt(x ** 2))\n",
    "    x = torch.randn(size=(10, 4))\n",
    "    x.requires_grad = True\n",
    "    assert gradient_penalty(lambda x: x, x, x).numpy() == 1\n",
    "    assert gradient_penalty(lambda x: x * 0, x, x).numpy() == 1\n",
    "\n",
    "\n",
    "test_gradient_penalty()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u9m30pey3eh6"
   },
   "source": [
    "Our training loop is similar to the previous one, but here we will \n",
    "use the gradient penalty instead of weight clipping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OH4WAoWXi1s0"
   },
   "outputs": [],
   "source": [
    "def train_wgan_gp(\n",
    "    generator, \n",
    "    critic, \n",
    "    train_loader,\n",
    "    critic_steps, \n",
    "    batch_size,\n",
    "    n_epochs,\n",
    "    lr, \n",
    "    gp_weight=10,\n",
    "    use_cuda=False\n",
    "):\n",
    "\n",
    "    if use_cuda:\n",
    "        critic = critic.cuda()\n",
    "        generator = generator.cuda()\n",
    "    critic.train()\n",
    "    generator.train()\n",
    "\n",
    "    gen_optimizer = torch.optim.Adam(generator.parameters(), lr=lr, betas=(0, 0.9))\n",
    "    critic_optimizer = torch.optim.Adam(critic.parameters(), lr=lr, betas=(0, 0.9))\n",
    "\n",
    "    curr_iter = 0\n",
    "    d_loss, g_loss = torch.zeros(1), torch.zeros(1)\n",
    "    batch_loss_history = {'discriminator_losses': [], 'generator_losses': []}\n",
    "    for epoch_i in tqdm(range(n_epochs)):\n",
    "        for batch_i, x in enumerate(train_loader):\n",
    "            curr_iter += 1\n",
    "            if use_cuda:\n",
    "                x = x.cuda()\n",
    "\n",
    "            # do a critic update\n",
    "            critic_optimizer.zero_grad()\n",
    "            fake_data = generator.sample(x.shape[0])\n",
    "\n",
    "            # ====\n",
    "            # your code\n",
    "            # D(x_fake) - D(x_real) + gp_weight * grad_pen\n",
    "\n",
    "            d_loss = \n",
    "            # ====\n",
    "\n",
    "            d_loss.backward()\n",
    "            critic_optimizer.step()\n",
    "            # generator update\n",
    "            if curr_iter % critic_steps == 0:\n",
    "                gen_optimizer.zero_grad()\n",
    "                fake_data = generator.sample(batch_size)\n",
    "                # ====\n",
    "                # your code\n",
    "                # -D(x_fake)\n",
    "                g_loss = \n",
    "                # ====\n",
    "                g_loss.backward()\n",
    "                gen_optimizer.step()\n",
    "\n",
    "                batch_loss_history['generator_losses'].append(g_loss.data.cpu().numpy())\n",
    "                batch_loss_history['discriminator_losses'].append(d_loss.data.cpu().numpy())\n",
    "\n",
    "    return batch_loss_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 731,
     "referenced_widgets": [
      "4636798c94b34e249cbb013049f8d8d9",
      "5fdd33344db14697af9a81e516c75dc4",
      "98d6add9f5e4485696094df5dfc37e2c",
      "a55ca40b7c8e4b148e764437af386d28",
      "542d9a0509694400a169e94d88f9fd7c",
      "a885e939ae32461cb98b44a664529f02",
      "306d12232e2f469abe140ee06cc85d7c",
      "81b25d863aef4d1191956666359b2f9c",
      "114da5b18ceb4b819626e99883baa325",
      "97e8832bcbcc46b389a01cb7ca6dea0f",
      "e87c9c512318455a9731fee6f7da938d"
     ]
    },
    "id": "M5RfNxTQjfU8",
    "outputId": "c9ffc34c-9024-42fe-f449-27e0c2b23e0b"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters\n",
    "BATCH_SIZE =     # any adequate value\n",
    "GEN_HIDDENS =    # 4 layers with < 128 neurons would be enough\n",
    "DISCR_HIDDENS =  # 4 layers with < 128 neurons would be enough \n",
    "CRITIC_STEPS =   # > 2\n",
    "LR =             # < 1e-2\n",
    "GP_WEIGHT =      # > 5\n",
    "# ====\n",
    "\n",
    "N_EPOCHS = 200\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "\n",
    "generator = MLPGenerator(input_dim=16, hiddens=GEN_HIDDENS, output_dim=2)\n",
    "critic = MLPCritic(input_dim=2, hiddens=DISCR_HIDDENS, output_dim=1)\n",
    "\n",
    "losses = train_wgan_gp(\n",
    "    generator, \n",
    "    critic, \n",
    "    train_loader,\n",
    "    CRITIC_STEPS, \n",
    "    batch_size=BATCH_SIZE, \n",
    "    n_epochs=N_EPOCHS,\n",
    "    lr=LR,\n",
    "    gp_weight=GP_WEIGHT,\n",
    "    use_cuda=USE_CUDA\n",
    ")\n",
    "\n",
    "plot_losses(train_losses['discriminator_losses'], 'Critic loss')\n",
    "plot_losses(train_losses['generator_losses'], 'Generator loss')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nDxPK7523vQM"
   },
   "source": [
    "Let plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 357
    },
    "id": "BLx-0sNajkID",
    "outputId": "40cf2740-b800-4407-99df-6bd11201e8e5"
   },
   "outputs": [],
   "source": [
    "samples, grid, critic_output = make_inference(generator, critic)\n",
    "visualize_critic_output(samples, train_data, grid, critic_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h8zInnwKIakS"
   },
   "source": [
    "# Task 3: WGAN-GP for CIFAR 10 (5pt)\n",
    "\n",
    "In this task you will fit [Wasserstein GAN](https://arxiv.org/abs/1701.07875) with [Gradient Penalty](https://arxiv.org/pdf/1704.00028.pdf) model to the CIFAR10 dataset.\n",
    "\n",
    "Download the data from [here](https://drive.google.com/file/d/16j3nrJV821VOkkuRz7aYam8TyIXLnNme/view?usp=sharing) (you could use the cell below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ykn5rCwf7sXZ",
    "outputId": "435adf8b-5c90-4d27-9d49-6acad0cf4836"
   },
   "outputs": [],
   "source": [
    "!gdown --id 16j3nrJV821VOkkuRz7aYam8TyIXLnNme"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 375
    },
    "id": "dAvSLXlKZ5q7",
    "outputId": "ff519324-7c46-4e9c-9dc3-4b4b26afdd25"
   },
   "outputs": [],
   "source": [
    "train_data, test_data = load_pickle('/content/cifar10.pkl')\n",
    "visualize_images(train_data, 'CIFAR10 samples')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nc9WGmMMkxTU"
   },
   "source": [
    "Here we will use convolution-based generator and critic.\n",
    "\n",
    "First of all, let define generator network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7a7fjcx0vL8R"
   },
   "outputs": [],
   "source": [
    "class ConvGenerator(nn.Module):\n",
    "    def __init__(self, input_size=128, n_channels=64):\n",
    "        super().__init__()\n",
    "        self.n_channels = n_channels\n",
    "        self.input_size = input_size\n",
    "        # ====\n",
    "        # your code\n",
    "        # you could you this architecture or your own\n",
    "        # 1) define linear layer with output units 4 * 4 * 4 * n_channels, then relu\n",
    "        # 2) define transposed conv with stride 2, kernel size 2 then BN, then relu\n",
    "        # 3) define transposed conv with stride 2, kernel size 2 then BN, then relu\n",
    "\n",
    "        # ====\n",
    "\n",
    "    def forward(self, input):\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) apply all layers\n",
    "        # 2) the output should be in the range of [0, 1] (apply activation) \n",
    "\n",
    "        # ====\n",
    "        return output.view(-1, 3, 32, 32)\n",
    "\n",
    "    def sample(self, n_samples):\n",
    "        # ====\n",
    "        # your code\n",
    "        # sample from standard normal distribution and apply the model\n",
    "\n",
    "        # ===="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7uVbwHvAruFh"
   },
   "source": [
    "Now it is time to define our critic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YFA0tI7ZrloP"
   },
   "outputs": [],
   "source": [
    "class ConvCritic(nn.Module):\n",
    "    def __init__(self, n_channels):\n",
    "        super().__init__()\n",
    "        self.n_channels = n_channels\n",
    "        \n",
    "        # ====\n",
    "        # your code\n",
    "        # sequence of Conv2D with kernel size 3 stride 2 and LeakyRelU\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def forward(self, x):\n",
    "        # ====\n",
    "        # your code\n",
    "        # apply all layers\n",
    "\n",
    "        # ====\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LBRZJXGAryNA"
   },
   "source": [
    "That is all :) We will use the train loop from the previous task here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 731,
     "referenced_widgets": [
      "9a5b6a8933124e5c9ca5e83e492f8161",
      "00b9b11c5866433fb1cdb10cd4d8e6b2",
      "a433ef3893c24c6b97bcc1795d850785",
      "3aa47e49d5294194a6e9b32fcb3372d4",
      "02dd7c4c820046e2888e437d32ee0de4",
      "b3c294ab69ad4ae8887f8492da928baa",
      "7789dd29012747c39da64c37f314faed",
      "bf5b7435a83e4e42bdd1356dccac2ea8",
      "915c9eac4d44435f94633c61868414ba",
      "5db49fbf8739420aa5f093349130ec8e",
      "540b43cee602436fbeacf2b4116e8709"
     ]
    },
    "id": "aY6qzJ5MSuiX",
    "outputId": "eb1b4122-5389-44bb-c211-2e953bbdebbc"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters (you have to train the model more than 20 epochs to get good results)\n",
    "BATCH_SIZE =    # any adequate value\n",
    "N_CHANNELS =    # > 32\n",
    "N_EPOCHS =      # > 10\n",
    "CRITIC_STEPS =  # > 2\n",
    "GP_WEIGHT =     # > 5\n",
    "LR =            # < 1e-3\n",
    "# ====\n",
    "\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "\n",
    "generator = ConvGenerator(n_channels=N_CHANNELS)\n",
    "critic = ConvCritic(n_channels=N_CHANNELS)\n",
    "\n",
    "train_losses = train_wgan_gp(\n",
    "    generator, \n",
    "    critic, \n",
    "    train_loader,\n",
    "    critic_steps=CRITIC_STEPS, \n",
    "    batch_size=BATCH_SIZE, \n",
    "    n_epochs=N_EPOCHS,\n",
    "    lr=LR,\n",
    "    gp_weight=GP_WEIGHT,\n",
    "    use_cuda=USE_CUDA\n",
    ")\n",
    "\n",
    "plot_losses(train_losses['discriminator_losses'], 'Discriminator loss')\n",
    "plot_losses(train_losses['generator_losses'], 'Generator loss')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "unAqs_pEs59l"
   },
   "source": [
    "Let sample from our model and draw the samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 375
    },
    "id": "8I0pNzHchqRs",
    "outputId": "739ec3d3-39aa-428b-87a0-28a71e56fa9a"
   },
   "outputs": [],
   "source": [
    "generator.eval()\n",
    "critic.eval()\n",
    "with torch.no_grad():\n",
    "    samples = generator.sample(1000)\n",
    "    samples = samples.cpu().detach().numpy()\n",
    "    \n",
    "\n",
    "show_samples(samples[:100], title='CIFAR-10 generated samples')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "hw5_solutions.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": {
    "00b9b11c5866433fb1cdb10cd4d8e6b2": {
     "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_b3c294ab69ad4ae8887f8492da928baa",
      "placeholder": "​",
      "style": "IPY_MODEL_7789dd29012747c39da64c37f314faed",
      "value": "100%"
     }
    },
    "02dd7c4c820046e2888e437d32ee0de4": {
     "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
     }
    },
    "114da5b18ceb4b819626e99883baa325": {
     "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": ""
     }
    },
    "14179864decd4f2390c42259e9bc426b": {
     "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
     }
    },
    "306d12232e2f469abe140ee06cc85d7c": {
     "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": ""
     }
    },
    "3aa47e49d5294194a6e9b32fcb3372d4": {
     "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_5db49fbf8739420aa5f093349130ec8e",
      "placeholder": "​",
      "style": "IPY_MODEL_540b43cee602436fbeacf2b4116e8709",
      "value": " 20/20 [05:22&lt;00:00, 16.44s/it]"
     }
    },
    "41585a0cbe524f4cbd03022e17393fe6": {
     "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": ""
     }
    },
    "4636798c94b34e249cbb013049f8d8d9": {
     "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_5fdd33344db14697af9a81e516c75dc4",
       "IPY_MODEL_98d6add9f5e4485696094df5dfc37e2c",
       "IPY_MODEL_a55ca40b7c8e4b148e764437af386d28"
      ],
      "layout": "IPY_MODEL_542d9a0509694400a169e94d88f9fd7c"
     }
    },
    "540b43cee602436fbeacf2b4116e8709": {
     "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": ""
     }
    },
    "542d9a0509694400a169e94d88f9fd7c": {
     "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
     }
    },
    "5b1cb14e8efa4cf5b7be757462f51d74": {
     "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
     }
    },
    "5db49fbf8739420aa5f093349130ec8e": {
     "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
     }
    },
    "5fdd33344db14697af9a81e516c75dc4": {
     "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_a885e939ae32461cb98b44a664529f02",
      "placeholder": "​",
      "style": "IPY_MODEL_306d12232e2f469abe140ee06cc85d7c",
      "value": "100%"
     }
    },
    "68765de1436245f192703b2d139ee0ac": {
     "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_5b1cb14e8efa4cf5b7be757462f51d74",
      "placeholder": "​",
      "style": "IPY_MODEL_8345ed2254f244f5bfdbf7bd5bf8c267",
      "value": "100%"
     }
    },
    "7789dd29012747c39da64c37f314faed": {
     "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": ""
     }
    },
    "81b25d863aef4d1191956666359b2f9c": {
     "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
     }
    },
    "8345ed2254f244f5bfdbf7bd5bf8c267": {
     "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": ""
     }
    },
    "84aa41c85adb4e11bfd484e3f316a839": {
     "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
     }
    },
    "915c9eac4d44435f94633c61868414ba": {
     "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": ""
     }
    },
    "97e8832bcbcc46b389a01cb7ca6dea0f": {
     "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
     }
    },
    "98d6add9f5e4485696094df5dfc37e2c": {
     "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_81b25d863aef4d1191956666359b2f9c",
      "max": 200,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_114da5b18ceb4b819626e99883baa325",
      "value": 200
     }
    },
    "9a51463240d84c21b23600b64cf1a096": {
     "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_14179864decd4f2390c42259e9bc426b",
      "placeholder": "​",
      "style": "IPY_MODEL_d72569fa2f5641fc8ca12367321944fc",
      "value": " 200/200 [00:35&lt;00:00,  6.72it/s]"
     }
    },
    "9a5b6a8933124e5c9ca5e83e492f8161": {
     "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_00b9b11c5866433fb1cdb10cd4d8e6b2",
       "IPY_MODEL_a433ef3893c24c6b97bcc1795d850785",
       "IPY_MODEL_3aa47e49d5294194a6e9b32fcb3372d4"
      ],
      "layout": "IPY_MODEL_02dd7c4c820046e2888e437d32ee0de4"
     }
    },
    "a433ef3893c24c6b97bcc1795d850785": {
     "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_bf5b7435a83e4e42bdd1356dccac2ea8",
      "max": 20,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_915c9eac4d44435f94633c61868414ba",
      "value": 20
     }
    },
    "a55ca40b7c8e4b148e764437af386d28": {
     "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_97e8832bcbcc46b389a01cb7ca6dea0f",
      "placeholder": "​",
      "style": "IPY_MODEL_e87c9c512318455a9731fee6f7da938d",
      "value": " 200/200 [00:39&lt;00:00,  5.27it/s]"
     }
    },
    "a885e939ae32461cb98b44a664529f02": {
     "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
     }
    },
    "ac937e59ba7f456ab73b2fcf1fe54014": {
     "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_84aa41c85adb4e11bfd484e3f316a839",
      "max": 200,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_41585a0cbe524f4cbd03022e17393fe6",
      "value": 200
     }
    },
    "b3c294ab69ad4ae8887f8492da928baa": {
     "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
     }
    },
    "bad267b7949340e1860b6475f81ab0de": {
     "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
     }
    },
    "bf5b7435a83e4e42bdd1356dccac2ea8": {
     "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
     }
    },
    "d72569fa2f5641fc8ca12367321944fc": {
     "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": ""
     }
    },
    "e87c9c512318455a9731fee6f7da938d": {
     "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": ""
     }
    },
    "effa9aab4edb4cc18261259a60f7d4e3": {
     "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_68765de1436245f192703b2d139ee0ac",
       "IPY_MODEL_ac937e59ba7f456ab73b2fcf1fe54014",
       "IPY_MODEL_9a51463240d84c21b23600b64cf1a096"
      ],
      "layout": "IPY_MODEL_bad267b7949340e1860b6475f81ab0de"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}