coursera-cpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1czVdIlqnImH"
   },
   "source": [
    "# Wasserstein GAN with Gradient Penalty (WGAN-GP)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1KD3ZgLs80vY"
   },
   "source": [
    "### Goals\n",
    "In this notebook, you're going to build a Wasserstein GAN with Gradient Penalty (WGAN-GP) that solves some of the stability issues with the GANs that you have been using up until this point. Specifically, you'll use a special kind of loss function known as the W-loss, where W stands for Wasserstein, and gradient penalties to prevent mode collapse.\n",
    "\n",
    "*Fun Fact: Wasserstein is named after a mathematician at Penn State, Leonid Vaseršteĭn. You'll see it abbreviated to W (e.g. WGAN, W-loss, W-distance).*\n",
    "\n",
    "### Learning Objectives\n",
    "1.   Get hands-on experience building a more stable GAN: Wasserstein GAN with Gradient Penalty (WGAN-GP).\n",
    "2.   Train the more advanced WGAN-GP model.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wU8DDM6l9rZb"
   },
   "source": [
    "## Generator and Critic\n",
    "\n",
    "You will begin by importing some useful packages, defining visualization functions, building the generator, and building the critic. Since the changes for WGAN-GP are done to the loss function during training, you can simply reuse your previous GAN code for the generator and critic class. Remember that in WGAN-GP, you no longer use a discriminator that classifies fake and real as 0 and 1 but rather a critic that scores images with real numbers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1sTpFE_eSk4s"
   },
   "source": [
    "#### Packages and Visualizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JfkorNJrnmNO"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from tqdm.auto import tqdm\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision.utils import make_grid\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "torch.manual_seed(0) # Set for testing purposes, please do not change!\n",
    "\n",
    "def show_tensor_images(image_tensor, num_images=25, size=(1, 28, 28)):\n",
    "    '''\n",
    "    Function for visualizing images: Given a tensor of images, number of images, and\n",
    "    size per image, plots and prints the images in an uniform grid.\n",
    "    '''\n",
    "    image_tensor = (image_tensor + 1) / 2\n",
    "    image_unflat = image_tensor.detach().cpu()\n",
    "    image_grid = make_grid(image_unflat[:num_images], nrow=5)\n",
    "    plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n",
    "    plt.show()\n",
    "\n",
    "def make_grad_hook():\n",
    "    '''\n",
    "    Function to keep track of gradients for visualization purposes, \n",
    "    which fills the grads list when using model.apply(grad_hook).\n",
    "    '''\n",
    "    grads = []\n",
    "    def grad_hook(m):\n",
    "        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n",
    "            grads.append(m.weight.grad)\n",
    "    return grads, grad_hook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P1A1M6kpnfxw"
   },
   "source": [
    "#### Generator and Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tFwajQ3tGgI2"
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    '''\n",
    "    Generator Class\n",
    "    Values:\n",
    "        z_dim: the dimension of the noise vector, a scalar\n",
    "        im_chan: the number of channels in the images, fitted for the dataset used, a scalar\n",
    "              (MNIST is black-and-white, so 1 channel is your default)\n",
    "        hidden_dim: the inner dimension, a scalar\n",
    "    '''\n",
    "    def __init__(self, z_dim=10, im_chan=1, hidden_dim=64):\n",
    "        super(Generator, self).__init__()\n",
    "        self.z_dim = z_dim\n",
    "        # Build the neural network\n",
    "        self.gen = nn.Sequential(\n",
    "            self.make_gen_block(z_dim, hidden_dim * 4),\n",
    "            self.make_gen_block(hidden_dim * 4, hidden_dim * 2, kernel_size=4, stride=1),\n",
    "            self.make_gen_block(hidden_dim * 2, hidden_dim),\n",
    "            self.make_gen_block(hidden_dim, im_chan, kernel_size=4, final_layer=True),\n",
    "        )\n",
    "\n",
    "    def make_gen_block(self, input_channels, output_channels, kernel_size=3, stride=2, final_layer=False):\n",
    "        '''\n",
    "        Function to return a sequence of operations corresponding to a generator block of DCGAN;\n",
    "        a transposed convolution, a batchnorm (except in the final layer), and an activation.\n",
    "        Parameters:\n",
    "            input_channels: how many channels the input feature representation has\n",
    "            output_channels: how many channels the output feature representation should have\n",
    "            kernel_size: the size of each convolutional filter, equivalent to (kernel_size, kernel_size)\n",
    "            stride: the stride of the convolution\n",
    "            final_layer: a boolean, true if it is the final layer and false otherwise \n",
    "                      (affects activation and batchnorm)\n",
    "        '''\n",
    "        if not final_layer:\n",
    "            return nn.Sequential(\n",
    "                nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),\n",
    "                nn.BatchNorm2d(output_channels),\n",
    "                nn.ReLU(inplace=True),\n",
    "            )\n",
    "        else:\n",
    "            return nn.Sequential(\n",
    "                nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),\n",
    "                nn.Tanh(),\n",
    "            )\n",
    "\n",
    "    def forward(self, noise):\n",
    "        '''\n",
    "        Function for completing a forward pass of the generator: Given a noise tensor,\n",
    "        returns generated images.\n",
    "        Parameters:\n",
    "            noise: a noise tensor with dimensions (n_samples, z_dim)\n",
    "        '''\n",
    "        x = noise.view(len(noise), self.z_dim, 1, 1)\n",
    "        print(x.size())\n",
    "        #return self.gen(x)\n",
    "        for layer in self.gen:\n",
    "            x = layer(x)\n",
    "            print(x.size())\n",
    "        print('Done')\n",
    "        return x\n",
    "\n",
    "def get_noise(n_samples, z_dim, device='cpu'):\n",
    "    '''\n",
    "    Function for creating noise vectors: Given the dimensions (n_samples, z_dim)\n",
    "    creates a tensor of that shape filled with random numbers from the normal distribution.\n",
    "    Parameters:\n",
    "      n_samples: the number of samples to generate, a scalar\n",
    "      z_dim: the dimension of the noise vector, a scalar\n",
    "      device: the device type\n",
    "    '''\n",
    "    return torch.randn(n_samples, z_dim, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "r9fScH98nkYH"
   },
   "source": [
    "#### Critic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aA4AxGnmpuPq"
   },
   "outputs": [],
   "source": [
    "class Critic(nn.Module):\n",
    "    '''\n",
    "    Critic Class\n",
    "    Values:\n",
    "        im_chan: the number of channels in the images, fitted for the dataset used, a scalar\n",
    "              (MNIST is black-and-white, so 1 channel is your default)\n",
    "        hidden_dim: the inner dimension, a scalar\n",
    "    '''\n",
    "    def __init__(self, im_chan=1, hidden_dim=64):\n",
    "        super(Critic, self).__init__()\n",
    "        self.crit = nn.Sequential(\n",
    "            self.make_crit_block(im_chan, hidden_dim),\n",
    "            self.make_crit_block(hidden_dim, hidden_dim * 2),\n",
    "            self.make_crit_block(hidden_dim * 2, 1, final_layer=True),\n",
    "        )\n",
    "\n",
    "    def make_crit_block(self, input_channels, output_channels, kernel_size=4, stride=2, final_layer=False):\n",
    "        '''\n",
    "        Function to return a sequence of operations corresponding to a critic block of DCGAN;\n",
    "        a convolution, a batchnorm (except in the final layer), and an activation (except in the final layer).\n",
    "        Parameters:\n",
    "            input_channels: how many channels the input feature representation has\n",
    "            output_channels: how many channels the output feature representation should have\n",
    "            kernel_size: the size of each convolutional filter, equivalent to (kernel_size, kernel_size)\n",
    "            stride: the stride of the convolution\n",
    "            final_layer: a boolean, true if it is the final layer and false otherwise \n",
    "                      (affects activation and batchnorm)\n",
    "        '''\n",
    "        if not final_layer:\n",
    "            return nn.Sequential(\n",
    "                nn.Conv2d(input_channels, output_channels, kernel_size, stride),\n",
    "                nn.BatchNorm2d(output_channels),\n",
    "                nn.LeakyReLU(0.2, inplace=True),\n",
    "            )\n",
    "        else:\n",
    "            return nn.Sequential(\n",
    "                nn.Conv2d(input_channels, output_channels, kernel_size, stride),\n",
    "            )\n",
    "\n",
    "    def forward(self, image):\n",
    "        '''\n",
    "        Function for completing a forward pass of the critic: Given an image tensor, \n",
    "        returns a 1-dimension tensor representing fake/real.\n",
    "        Parameters:\n",
    "            image: a flattened image tensor with dimension (im_chan)\n",
    "        '''\n",
    "        crit_pred = self.crit(image)\n",
    "        return crit_pred.view(len(crit_pred), -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qRk_8azSq3tF"
   },
   "source": [
    "## Training Initializations\n",
    "Now you can start putting it all together.\n",
    "As usual, you will start by setting the parameters:\n",
    "  *   n_epochs: the number of times you iterate through the entire dataset when training\n",
    "  *   z_dim: the dimension of the noise vector\n",
    "  *   display_step: how often to display/visualize the images\n",
    "  *   batch_size: the number of images per forward/backward pass\n",
    "  *   lr: the learning rate\n",
    "  *   beta_1, beta_2: the momentum terms\n",
    "  *   c_lambda: weight of the gradient penalty\n",
    "  *   crit_repeats: number of times to update the critic per generator update - there are more details about this in the *Putting It All Together* section\n",
    "  *   device: the device type\n",
    "\n",
    "You will also load and transform the MNIST dataset to tensors.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 352,
     "referenced_widgets": [
      "b16f592cc3964cccb48ea8247a83d93c",
      "1d0d36a15ed24a11a64631880986cbe2",
      "b7291a08accb4c978c9c6f368e0958bb",
      "b8115cdd58e64d179e683903ad468a66",
      "25ac91c871744f10b4c6e4f1099fecfa",
      "ba6a57ef686a4e128e0ac5b61ec7bbb0",
      "3f8115b26eb744a9a66dd5952a76e486",
      "3a44d4d196c940e186f2ad7c978af58e",
      "c02e809ef8bd41c39c6685fc9ce9e97e",
      "be0abdaac7744d3a915ea6f710f482f0",
      "b426f0363f5e4bbfb442b16b1482e98a",
      "2abda2d9d9ca4970b528e2e29f12f4b4",
      "c861f7f9c6db4e60b9e9cc3571c1d558",
      "1fb9f3f27a7d4784b944479f8724ce8b",
      "61d57a6747f542779a1fa9903d63b270",
      "a34a69313e64475bbea18f55d3d570c7",
      "bd82c2e94e224c1ca71e9bfee2facf16",
      "b11c1e9c1c864c87ade7897e66108551",
      "f3138380792542a3aa7a5ee7fd8fca44",
      "46356cd2605a4f97b60e774df34f4e4b",
      "44b8f12ebf6d4928b8ec5939d2a4f43f",
      "90e46596d5034e058638960da71cd3de",
      "f043c3530998460eb8a540f673994e4e",
      "4311158e895448ebb68efe079896f4b5",
      "9041af49322f4246af43f1ed7579063b",
      "d9c11e9a1d6e4327bde7c917bc0db171",
      "b503a11497f64a63961287665024a499",
      "7d5c4bbce6ec4e6bb61a415ad968342d",
      "bdb5a2d8e9ac4f9db45a6c4002b22c2d",
      "ce2757c985e14276bf7222c020c473a4",
      "0c164c68cf6849f89194d2d5e0ed4b07",
      "252af3366dab40068dbbe9845f098560"
     ]
    },
    "colab_type": "code",
    "id": "IFLQ039u-qdu",
    "outputId": "2969e573-0b53-49e0-b1e9-ac6058d5d6b2"
   },
   "outputs": [],
   "source": [
    "n_epochs = 100\n",
    "z_dim = 64\n",
    "display_step = 500\n",
    "batch_size = 128\n",
    "lr = 0.0002\n",
    "beta_1 = 0.5\n",
    "beta_2 = 0.999\n",
    "c_lambda = 10\n",
    "crit_repeats = 5\n",
    "device = 'cpu'\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,)),\n",
    "])\n",
    "\n",
    "dataloader = DataLoader(\n",
    "    MNIST('.', download=True, transform=transform),\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "24Var22i_Ccs"
   },
   "source": [
    "Then, you can initialize your generator, critic, and optimizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "sDFRZ8tg_Y57",
    "outputId": "d759a7d3-48e9-450b-81a0-dc5cfe406a8d"
   },
   "outputs": [],
   "source": [
    "gen = Generator(z_dim).to(device)\n",
    "gen_opt = torch.optim.Adam(gen.parameters(), lr=lr, betas=(beta_1, beta_2))\n",
    "crit = Critic().to(device) \n",
    "crit_opt = torch.optim.Adam(crit.parameters(), lr=lr, betas=(beta_1, beta_2))\n",
    "\n",
    "def weights_init(m):\n",
    "    if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n",
    "        torch.nn.init.normal_(m.weight, 0.0, 0.02)\n",
    "    if isinstance(m, nn.BatchNorm2d):\n",
    "        torch.nn.init.normal_(m.weight, 0.0, 0.02)\n",
    "        torch.nn.init.constant_(m.bias, 0)\n",
    "gen = gen.apply(weights_init)\n",
    "crit = crit.apply(weights_init)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BFEi5BhVX5-P"
   },
   "source": [
    "## Gradient Penalty\n",
    "Calculating the gradient penalty can be broken into two functions: (1) compute the gradient with respect to the images and (2) compute the gradient penalty given the gradient.\n",
    "\n",
    "You can start by getting the gradient. The gradient is computed by first creating a mixed image. This is done by weighing the fake and real image using epsilon and then adding them together. Once you have the intermediate image, you can get the critic's output on the image. Finally, you compute the gradient of the critic score's on the mixed images (output) with respect to the pixels of the mixed images (input). You will need to fill in the code to get the gradient wherever you see *None*. There is a test function in the next block for you to test your solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Tn4dkXnNtcv6"
   },
   "outputs": [],
   "source": [
    "# UNQ_C1 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: get_gradient\n",
    "def get_gradient(crit, real, fake, epsilon):\n",
    "    '''\n",
    "    Return the gradient of the critic's scores with respect to mixes of real and fake images.\n",
    "    Parameters:\n",
    "        crit: the critic model\n",
    "        real: a batch of real images\n",
    "        fake: a batch of fake images\n",
    "        epsilon: a vector of the uniformly random proportions of real/fake per mixed image\n",
    "    Returns:\n",
    "        gradient: the gradient of the critic's scores, with respect to the mixed image\n",
    "    '''\n",
    "    # Mix the images together\n",
    "    mixed_images = real * epsilon + fake * (1 - epsilon)\n",
    "\n",
    "    # Calculate the critic's scores on the mixed images\n",
    "    mixed_scores = crit(mixed_images)\n",
    "    \n",
    "    # Take the gradient of the scores with respect to the images\n",
    "    gradient = torch.autograd.grad(\n",
    "        # Note: You need to take the gradient of outputs with respect to inputs.\n",
    "        # This documentation may be useful, but it should not be necessary:\n",
    "        # https://pytorch.org/docs/stable/autograd.html#torch.autograd.grad\n",
    "        #### START CODE HERE ####\n",
    "        inputs=mixed_images,\n",
    "        outputs=mixed_scores,\n",
    "        #### END CODE HERE ####\n",
    "        # These other parameters have to do with the pytorch autograd engine works\n",
    "        grad_outputs=torch.ones_like(mixed_scores), \n",
    "        create_graph=True,\n",
    "        retain_graph=True,\n",
    "    )[0]\n",
    "    return gradient\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8av4TtbMtkTq"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success!\n"
     ]
    }
   ],
   "source": [
    "# UNIT TEST\n",
    "# DO NOT MODIFY THIS\n",
    "def test_get_gradient(image_shape):\n",
    "    real = torch.randn(*image_shape, device=device) + 1\n",
    "    fake = torch.randn(*image_shape, device=device) - 1\n",
    "    epsilon_shape = [1 for _ in image_shape]\n",
    "    epsilon_shape[0] = image_shape[0]\n",
    "    epsilon = torch.rand(epsilon_shape, device=device).requires_grad_()\n",
    "    gradient = get_gradient(crit, real, fake, epsilon)\n",
    "    assert tuple(gradient.shape) == image_shape\n",
    "    assert gradient.max() > 0\n",
    "    assert gradient.min() < 0\n",
    "    return gradient\n",
    "\n",
    "gradient = test_get_gradient((256, 1, 28, 28))\n",
    "print(\"Success!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Z5PMRrMpRUK-"
   },
   "source": [
    "The second function you need to complete is to compute the gradient penalty given the gradient. First, you calculate the magnitude of each image's gradient. The magnitude of a gradient is also called the norm. Then, you calculate the penalty by squaring the distance between each magnitude and the ideal norm of 1 and taking the mean of all the squared distances.\n",
    "\n",
    "Again, you will need to fill in the code wherever you see *None*. There are hints below that you can view if you need help and there is a test function in the next block for you to test your solution.\n",
    "\n",
    "<details>\n",
    "\n",
    "<summary>\n",
    "<font size=\"3\" color=\"green\">\n",
    "<b>Optional hints for <code><font size=\"4\">gradient_penalty</font></code></b>\n",
    "</font>\n",
    "</summary>\n",
    "\n",
    "\n",
    "1.   Make sure you take the mean at the end.\n",
    "2.   Note that the magnitude of each gradient has already been calculated for you.\n",
    "\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VPwBH83IzCpS"
   },
   "outputs": [],
   "source": [
    "# UNQ_C2 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: gradient_penalty\n",
    "def gradient_penalty(gradient):\n",
    "    '''\n",
    "    Return the gradient penalty, given a gradient.\n",
    "    Given a batch of image gradients, you calculate the magnitude of each image's gradient\n",
    "    and penalize the mean quadratic distance of each magnitude to 1.\n",
    "    Parameters:\n",
    "        gradient: the gradient of the critic's scores, with respect to the mixed image\n",
    "    Returns:\n",
    "        penalty: the gradient penalty\n",
    "    '''\n",
    "    # Flatten the gradients so that each row captures one image\n",
    "    gradient = gradient.view(len(gradient), -1)\n",
    "\n",
    "    # Calculate the magnitude of every row\n",
    "    gradient_norm = gradient.norm(2, dim=1)\n",
    "    \n",
    "    # Penalize the mean squared distance of the gradient norms from 1\n",
    "    #### START CODE HERE ####\n",
    "    penalty = torch.mean((gradient_norm - 1) ** 2)\n",
    "    #### END CODE HERE ####\n",
    "    return penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ahPfGMA2zABQ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success!\n"
     ]
    }
   ],
   "source": [
    "# UNIT TEST\n",
    "def test_gradient_penalty(image_shape):\n",
    "    bad_gradient = torch.zeros(*image_shape)\n",
    "    bad_gradient_penalty = gradient_penalty(bad_gradient)\n",
    "    assert torch.isclose(bad_gradient_penalty, torch.tensor(1.))\n",
    "\n",
    "    image_size = torch.prod(torch.Tensor(image_shape[1:]))\n",
    "    good_gradient = torch.ones(*image_shape) / torch.sqrt(image_size)\n",
    "    good_gradient_penalty = gradient_penalty(good_gradient)\n",
    "    assert torch.isclose(good_gradient_penalty, torch.tensor(0.))\n",
    "\n",
    "    random_gradient = test_get_gradient(image_shape)\n",
    "    random_gradient_penalty = gradient_penalty(random_gradient)\n",
    "    assert torch.abs(random_gradient_penalty - 1) < 0.1\n",
    "\n",
    "test_gradient_penalty((256, 1, 28, 28))\n",
    "print(\"Success!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Sob-u9Z1X9sb"
   },
   "source": [
    "## Losses\n",
    "Next, you need to calculate the loss for the generator and the critic.\n",
    "\n",
    "For the generator, the loss is calculated by maximizing the critic's prediction on the generator's fake images. The argument has the scores for all fake images in the batch, but you will use the mean of them.\n",
    "\n",
    "There are optional hints below and a test function in the next block for you to test your solution.\n",
    "\n",
    "<details><summary><font size=\"3\" color=\"green\"><b>Optional hints for <code><font size=\"4\">get_gen_loss</font></code></b></font></summary>\n",
    "\n",
    "1. This can be written in one line.\n",
    "2. This is the negative of the mean of the critic's scores.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YnJFs-qkMCA-"
   },
   "outputs": [],
   "source": [
    "# UNQ_C3 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: get_gen_loss\n",
    "def get_gen_loss(crit_fake_pred):\n",
    "    '''\n",
    "    Return the loss of a generator given the critic's scores of the generator's fake images.\n",
    "    Parameters:\n",
    "        crit_fake_pred: the critic's scores of the fake images\n",
    "    Returns:\n",
    "        gen_loss: a scalar loss value for the current batch of the generator\n",
    "    '''\n",
    "    #### START CODE HERE ####\n",
    "    gen_loss = - torch.mean(crit_fake_pred)\n",
    "    #### END CODE HERE ####\n",
    "    return gen_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fYVqG8bR6Hfg"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success!\n"
     ]
    }
   ],
   "source": [
    "# UNIT TEST\n",
    "assert torch.isclose(\n",
    "    get_gen_loss(torch.tensor(1.)), torch.tensor(-1.0)\n",
    ")\n",
    "\n",
    "assert torch.isclose(\n",
    "    get_gen_loss(torch.rand(10000)), torch.tensor(-0.5), 0.05\n",
    ")\n",
    "\n",
    "print(\"Success!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vt3MgH68TM1P"
   },
   "source": [
    "For the critic, the loss is calculated by maximizing the distance between the critic's predictions on the real images and the predictions on the fake images while also adding a gradient penalty. The gradient penalty is weighed according to lambda. The arguments are the scores for all the images in the batch, and you will use the mean of them.\n",
    "\n",
    "There are hints below if you get stuck and a test function in the next block for you to test your solution.\n",
    "\n",
    "<details><summary><font size=\"3\" color=\"green\"><b>Optional hints for <code><font size=\"4\">get_crit_loss</font></code></b></font></summary>\n",
    "\n",
    "1. The higher the mean fake score, the higher the critic's loss is.\n",
    "2. What does this suggest about the mean real score?\n",
    "3. The higher the gradient penalty, the higher the critic's loss is, proportional to lambda.\n",
    "\n",
    "\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-jvbz1zDMTdu"
   },
   "outputs": [],
   "source": [
    "# UNQ_C4 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: get_crit_loss\n",
    "def get_crit_loss(crit_fake_pred, crit_real_pred, gp, c_lambda):\n",
    "    '''\n",
    "    Return the loss of a critic given the critic's scores for fake and real images,\n",
    "    the gradient penalty, and gradient penalty weight.\n",
    "    Parameters:\n",
    "        crit_fake_pred: the critic's scores of the fake images\n",
    "        crit_real_pred: the critic's scores of the real images\n",
    "        gp: the unweighted gradient penalty\n",
    "        c_lambda: the current weight of the gradient penalty \n",
    "    Returns:\n",
    "        crit_loss: a scalar for the critic's loss, accounting for the relevant factors\n",
    "    '''\n",
    "    #### START CODE HERE ####\n",
    "    crit_loss = torch.mean(crit_fake_pred) - torch.mean(crit_real_pred) + c_lambda * gp\n",
    "    #### END CODE HERE ####\n",
    "    return crit_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dxZey6fc5luf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success!\n"
     ]
    }
   ],
   "source": [
    "# UNIT TEST\n",
    "assert torch.isclose(\n",
    "    get_crit_loss(torch.tensor(1.), torch.tensor(2.), torch.tensor(3.), 0.1),\n",
    "    torch.tensor(-0.7)\n",
    ")\n",
    "assert torch.isclose(\n",
    "    get_crit_loss(torch.tensor(20.), torch.tensor(-20.), torch.tensor(2.), 10),\n",
    "    torch.tensor(60.)\n",
    ")\n",
    "\n",
    "print(\"Success!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_x5wu7rUMlnZ"
   },
   "source": [
    "## Putting It All Together\n",
    "Before you put everything together, there are a few things to note.\n",
    "1.   Even on GPU, the **training will run more slowly** than previous labs because the gradient penalty requires you to compute the gradient of a gradient -- this means potentially a few minutes per epoch! For best results, run this for as long as you can while on GPU.\n",
    "2.   One important difference from earlier versions is that you will **update the critic multiple times** every time you update the generator This helps prevent the generator from overpowering the critic. Sometimes, you might see the reverse, with the generator updated more times than the critic. This depends on architectural (e.g. the depth and width of the network) and algorithmic choices (e.g. which loss you're using). \n",
    "3.   WGAN-GP isn't necessarily meant to improve overall performance of a GAN, but just **increases stability** and avoids mode collapse. In general, a WGAN will be able to train in a much more stable way than the vanilla DCGAN from last assignment, though it will generally run a bit slower. You should also be able to train your model for more epochs without it collapsing.\n",
    "\n",
    "\n",
    "<!-- Once again, be warned that this runs very slowly on a CPU. One way to run this more quickly is to download the .ipynb and upload it to Google Drive, then open it with Google Colab and make the runtime type GPU and replace\n",
    "`device = \"cpu\"`\n",
    "with\n",
    "`device = \"cuda\"`\n",
    "and make sure that your `get_noise` function uses the right device.  -->\n",
    "\n",
    "Here is a snapshot of what your WGAN-GP outputs should resemble:\n",
    "![MNIST Digits Progression](MNIST_WGAN_Progression.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UXptQZcwrBrq"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2c4c8da43524b95a86d5d4c0fb57ae3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=469.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([128, 64, 1, 1])\n",
      "torch.Size([128, 256, 3, 3])\n",
      "torch.Size([128, 128, 6, 6])\n",
      "torch.Size([128, 64, 13, 13])\n",
      "torch.Size([128, 1, 28, 28])\n",
      "Done\n",
      "torch.Size([128, 64, 1, 1])\n",
      "torch.Size([128, 256, 3, 3])\n",
      "torch.Size([128, 128, 6, 6])\n",
      "torch.Size([128, 64, 13, 13])\n",
      "torch.Size([128, 1, 28, 28])\n",
      "Done\n",
      "torch.Size([128, 64, 1, 1])\n",
      "torch.Size([128, 256, 3, 3])\n",
      "torch.Size([128, 128, 6, 6])\n",
      "torch.Size([128, 64, 13, 13])\n",
      "torch.Size([128, 1, 28, 28])\n",
      "Done\n",
      "torch.Size([128, 64, 1, 1])\n",
      "torch.Size([128, 256, 3, 3])\n",
      "torch.Size([128, 128, 6, 6])\n",
      "torch.Size([128, 64, 13, 13])\n",
      "torch.Size([128, 1, 28, 28])\n",
      "Done\n",
      "\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-29-8edc9aa1c6ab>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m             \u001b[0mepsilon\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrand\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreal\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrequires_grad\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m             \u001b[0mgradient\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_gradient\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcrit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfake\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepsilon\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m             \u001b[0mgp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgradient_penalty\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgradient\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m             \u001b[0mcrit_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_crit_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcrit_fake_pred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcrit_real_pred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mc_lambda\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-21-35150bdb454c>\u001b[0m in \u001b[0;36mget_gradient\u001b[0;34m(crit, real, fake, epsilon)\u001b[0m\n\u001b[1;32m     30\u001b[0m         \u001b[0mgrad_outputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones_like\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmixed_scores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     31\u001b[0m         \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 32\u001b[0;31m         \u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     33\u001b[0m     )[0]\n\u001b[1;32m     34\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mgrad\u001b[0;34m(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)\u001b[0m\n\u001b[1;32m    202\u001b[0m     return Variable._execution_engine.run_backward(\n\u001b[1;32m    203\u001b[0m         \u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_outputs_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 204\u001b[0;31m         inputs, allow_unused)\n\u001b[0m\u001b[1;32m    205\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    206\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "cur_step = 0\n",
    "generator_losses = []\n",
    "critic_losses = []\n",
    "for epoch in range(n_epochs):\n",
    "    # Dataloader returns the batches\n",
    "    for real, _ in tqdm(dataloader):\n",
    "        cur_batch_size = len(real)\n",
    "        real = real.to(device)\n",
    "\n",
    "        mean_iteration_critic_loss = 0\n",
    "        for _ in range(crit_repeats):\n",
    "            ### Update critic ###\n",
    "            crit_opt.zero_grad()\n",
    "            fake_noise = get_noise(cur_batch_size, z_dim, device=device)\n",
    "            fake = gen(fake_noise)\n",
    "            crit_fake_pred = crit(fake.detach())\n",
    "            crit_real_pred = crit(real)\n",
    "\n",
    "            epsilon = torch.rand(len(real), 1, 1, 1, device=device, requires_grad=True)\n",
    "            gradient = get_gradient(crit, real, fake.detach(), epsilon)\n",
    "            gp = gradient_penalty(gradient)\n",
    "            crit_loss = get_crit_loss(crit_fake_pred, crit_real_pred, gp, c_lambda)\n",
    "\n",
    "            # Keep track of the average critic loss in this batch\n",
    "            mean_iteration_critic_loss += crit_loss.item() / crit_repeats\n",
    "            # Update gradients\n",
    "            crit_loss.backward(retain_graph=True)\n",
    "            # Update optimizer\n",
    "            crit_opt.step()\n",
    "        critic_losses += [mean_iteration_critic_loss]\n",
    "\n",
    "        ### Update generator ###\n",
    "        gen_opt.zero_grad()\n",
    "        fake_noise_2 = get_noise(cur_batch_size, z_dim, device=device)\n",
    "        fake_2 = gen(fake_noise_2)\n",
    "        crit_fake_pred = crit(fake_2)\n",
    "        \n",
    "        gen_loss = get_gen_loss(crit_fake_pred)\n",
    "        gen_loss.backward()\n",
    "\n",
    "        # Update the weights\n",
    "        gen_opt.step()\n",
    "\n",
    "        # Keep track of the average generator loss\n",
    "        generator_losses += [gen_loss.item()]\n",
    "\n",
    "        ### Visualization code ###\n",
    "        if cur_step % display_step == 0 and cur_step > 0:\n",
    "            gen_mean = sum(generator_losses[-display_step:]) / display_step\n",
    "            crit_mean = sum(critic_losses[-display_step:]) / display_step\n",
    "#            print(f\"Step {cur_step}: Generator loss: {gen_mean}, critic loss: {crit_mean}\")\n",
    "            show_tensor_images(fake)\n",
    "            show_tensor_images(real)\n",
    "            step_bins = 20\n",
    "            num_examples = (len(generator_losses) // step_bins) * step_bins\n",
    "            plt.plot(\n",
    "                range(num_examples // step_bins), \n",
    "                torch.Tensor(generator_losses[:num_examples]).view(-1, step_bins).mean(1),\n",
    "                label=\"Generator Loss\"\n",
    "            )\n",
    "            plt.plot(\n",
    "                range(num_examples // step_bins), \n",
    "                torch.Tensor(critic_losses[:num_examples]).view(-1, step_bins).mean(1),\n",
    "                label=\"Critic Loss\"\n",
    "            )\n",
    "            plt.legend()\n",
    "            plt.show()\n",
    "\n",
    "        cur_step += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "1sTpFE_eSk4s"
   ],
   "name": "C1W3_3: WGAN-GP.ipynb",
   "provenance": []
  },
  "coursera": {
   "schema_names": [
    "GANSC1-3A"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.5.3"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0c164c68cf6849f89194d2d5e0ed4b07": {
     "model_module": "@jupyter-widgets/controls",
     "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": ""
     }
    },
    "1d0d36a15ed24a11a64631880986cbe2": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "1fb9f3f27a7d4784b944479f8724ce8b": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "252af3366dab40068dbbe9845f098560": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "25ac91c871744f10b4c6e4f1099fecfa": {
     "model_module": "@jupyter-widgets/controls",
     "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": "initial"
     }
    },
    "2abda2d9d9ca4970b528e2e29f12f4b4": {
     "model_module": "@jupyter-widgets/controls",
     "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_a34a69313e64475bbea18f55d3d570c7",
      "placeholder": "​",
      "style": "IPY_MODEL_61d57a6747f542779a1fa9903d63b270",
      "value": " 0/28881 [00:00&lt;?, ?it/s]"
     }
    },
    "3a44d4d196c940e186f2ad7c978af58e": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "3f8115b26eb744a9a66dd5952a76e486": {
     "model_module": "@jupyter-widgets/controls",
     "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": ""
     }
    },
    "4311158e895448ebb68efe079896f4b5": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "44b8f12ebf6d4928b8ec5939d2a4f43f": {
     "model_module": "@jupyter-widgets/controls",
     "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": "initial"
     }
    },
    "46356cd2605a4f97b60e774df34f4e4b": {
     "model_module": "@jupyter-widgets/controls",
     "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_4311158e895448ebb68efe079896f4b5",
      "placeholder": "​",
      "style": "IPY_MODEL_f043c3530998460eb8a540f673994e4e",
      "value": " 1654784/? [00:04&lt;00:00, 388186.87it/s]"
     }
    },
    "61d57a6747f542779a1fa9903d63b270": {
     "model_module": "@jupyter-widgets/controls",
     "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": ""
     }
    },
    "7d5c4bbce6ec4e6bb61a415ad968342d": {
     "model_module": "@jupyter-widgets/controls",
     "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_252af3366dab40068dbbe9845f098560",
      "placeholder": "​",
      "style": "IPY_MODEL_0c164c68cf6849f89194d2d5e0ed4b07",
      "value": " 8192/? [00:03&lt;00:00, 2363.67it/s]"
     }
    },
    "9041af49322f4246af43f1ed7579063b": {
     "model_module": "@jupyter-widgets/controls",
     "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_b503a11497f64a63961287665024a499",
       "IPY_MODEL_7d5c4bbce6ec4e6bb61a415ad968342d"
      ],
      "layout": "IPY_MODEL_d9c11e9a1d6e4327bde7c917bc0db171"
     }
    },
    "90e46596d5034e058638960da71cd3de": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "a34a69313e64475bbea18f55d3d570c7": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "b11c1e9c1c864c87ade7897e66108551": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "b16f592cc3964cccb48ea8247a83d93c": {
     "model_module": "@jupyter-widgets/controls",
     "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_b7291a08accb4c978c9c6f368e0958bb",
       "IPY_MODEL_b8115cdd58e64d179e683903ad468a66"
      ],
      "layout": "IPY_MODEL_1d0d36a15ed24a11a64631880986cbe2"
     }
    },
    "b426f0363f5e4bbfb442b16b1482e98a": {
     "model_module": "@jupyter-widgets/controls",
     "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": "info",
      "description": "  0%",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1fb9f3f27a7d4784b944479f8724ce8b",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_c861f7f9c6db4e60b9e9cc3571c1d558",
      "value": 0
     }
    },
    "b503a11497f64a63961287665024a499": {
     "model_module": "@jupyter-widgets/controls",
     "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_ce2757c985e14276bf7222c020c473a4",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_bdb5a2d8e9ac4f9db45a6c4002b22c2d",
      "value": 1
     }
    },
    "b7291a08accb4c978c9c6f368e0958bb": {
     "model_module": "@jupyter-widgets/controls",
     "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": "info",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ba6a57ef686a4e128e0ac5b61ec7bbb0",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_25ac91c871744f10b4c6e4f1099fecfa",
      "value": 1
     }
    },
    "b8115cdd58e64d179e683903ad468a66": {
     "model_module": "@jupyter-widgets/controls",
     "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_3a44d4d196c940e186f2ad7c978af58e",
      "placeholder": "​",
      "style": "IPY_MODEL_3f8115b26eb744a9a66dd5952a76e486",
      "value": " 9920512/? [00:20&lt;00:00, 748579.20it/s]"
     }
    },
    "ba6a57ef686a4e128e0ac5b61ec7bbb0": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "bd82c2e94e224c1ca71e9bfee2facf16": {
     "model_module": "@jupyter-widgets/controls",
     "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_f3138380792542a3aa7a5ee7fd8fca44",
       "IPY_MODEL_46356cd2605a4f97b60e774df34f4e4b"
      ],
      "layout": "IPY_MODEL_b11c1e9c1c864c87ade7897e66108551"
     }
    },
    "bdb5a2d8e9ac4f9db45a6c4002b22c2d": {
     "model_module": "@jupyter-widgets/controls",
     "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": "initial"
     }
    },
    "be0abdaac7744d3a915ea6f710f482f0": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "c02e809ef8bd41c39c6685fc9ce9e97e": {
     "model_module": "@jupyter-widgets/controls",
     "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_b426f0363f5e4bbfb442b16b1482e98a",
       "IPY_MODEL_2abda2d9d9ca4970b528e2e29f12f4b4"
      ],
      "layout": "IPY_MODEL_be0abdaac7744d3a915ea6f710f482f0"
     }
    },
    "c861f7f9c6db4e60b9e9cc3571c1d558": {
     "model_module": "@jupyter-widgets/controls",
     "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": "initial"
     }
    },
    "ce2757c985e14276bf7222c020c473a4": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "d9c11e9a1d6e4327bde7c917bc0db171": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "f043c3530998460eb8a540f673994e4e": {
     "model_module": "@jupyter-widgets/controls",
     "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": ""
     }
    },
    "f3138380792542a3aa7a5ee7fd8fca44": {
     "model_module": "@jupyter-widgets/controls",
     "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_90e46596d5034e058638960da71cd3de",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_44b8f12ebf6d4928b8ec5939d2a4f43f",
      "value": 1
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}