MoG_Exp1_2.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KX7W22OEMcC4",
        "outputId": "70ad96ba-caf4-4f97-fb3a-6e54fa53e586"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "64X938ISuDpJ"
      },
      "outputs": [],
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "import torch.optim as optim\n",
        "from torch import autograd\n",
        "\n",
        "import numpy as np\n",
        "import scipy as sp\n",
        "from scipy import stats\n",
        "from matplotlib import pyplot as plt\n",
        "import sys, os\n",
        "import copy\n",
        "import tqdm\n",
        "\n",
        "import math\n",
        "from torch.optim import Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0K5ZEuEv5gvb"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "path = '/content/drive/MyDrive/Numgan'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "usPVqBZjMv_H"
      },
      "outputs": [],
      "source": [
        "## Extragradient Algo\n",
        "\n",
        "\n",
        "required = object()\n",
        "\n",
        "class Extragradient(Optimizer):\n",
        "    \"\"\"Base class for optimizers with extrapolation step.\n",
        "        Arguments:\n",
        "        params (iterable): an iterable of :class:`torch.Tensor` s or\n",
        "            :class:`dict` s. Specifies what Tensors should be optimized.\n",
        "        defaults: (dict): a dict containing default values of optimization\n",
        "            options (used when a parameter group doesn't specify them).\n",
        "    \"\"\"\n",
        "    def __init__(self, params, defaults):\n",
        "        super(Extragradient, self).__init__(params, defaults)\n",
        "        self.params_copy = []\n",
        "\n",
        "    def update(self, p, group):\n",
        "        raise NotImplementedError\n",
        "\n",
        "    def extrapolation(self):\n",
        "        \"\"\"Performs the extrapolation step and save a copy of the current parameters for the update step.\n",
        "        \"\"\"\n",
        "        # Check if a copy of the parameters was already made.\n",
        "        is_empty = len(self.params_copy) == 0\n",
        "        for group in self.param_groups:\n",
        "            for p in group['params']:\n",
        "                u = self.update(p, group)\n",
        "                if is_empty:\n",
        "                    # Save the current parameters for the update step. Several extrapolation step can be made before each update but only the parameters before the first extrapolation step are saved.\n",
        "                    self.params_copy.append(p.data.clone())\n",
        "                if u is None:\n",
        "                    continue\n",
        "                # Update the current parameters\n",
        "                p.data.add_(u)\n",
        "\n",
        "    def step(self, closure=None):\n",
        "        \"\"\"Performs a single optimization step.\n",
        "        Arguments:\n",
        "            closure (callable, optional): A closure that reevaluates the model\n",
        "                and returns the loss.\n",
        "        \"\"\"\n",
        "        if len(self.params_copy) == 0:\n",
        "            raise RuntimeError('Need to call extrapolation before calling step.')\n",
        "\n",
        "        loss = None\n",
        "        if closure is not None:\n",
        "            loss = closure()\n",
        "\n",
        "        i = -1\n",
        "        for group in self.param_groups:\n",
        "            for p in group['params']:\n",
        "                i += 1\n",
        "                u = self.update(p, group)\n",
        "                if u is None:\n",
        "                    continue\n",
        "                # Update the parameters saved during the extrapolation step\n",
        "                p.data = self.params_copy[i].add_(u)\n",
        "\n",
        "\n",
        "        # Free the old parameters\n",
        "        self.params_copy = []\n",
        "        return loss\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "class ExtraAdam(Extragradient):\n",
        "    \"\"\"Implements the Adam algorithm with extrapolation step.\n",
        "    Arguments:\n",
        "        params (iterable): iterable of parameters to optimize or dicts defining\n",
        "            parameter groups\n",
        "        lr (float, optional): learning rate (default: 1e-3)\n",
        "        betas (Tuple[float, float], optional): coefficients used for computing\n",
        "            running averages of gradient and its square (default: (0.9, 0.999))\n",
        "        eps (float, optional): term added to the denominator to improve\n",
        "            numerical stability (default: 1e-8)\n",
        "        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n",
        "        amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n",
        "            algorithm from the paper `On the Convergence of Adam and Beyond`_\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n",
        "                 weight_decay=0, amsgrad=False):\n",
        "        if not 0.0 <= lr:\n",
        "         raise ValueError(\"Invalid learning rate: {}\".format(lr))\n",
        "        if not 0.0 <= eps:\n",
        "         raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n",
        "        if not 0.0 <= betas[0] < 1.0:\n",
        "         raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n",
        "        if not 0.0 <= betas[1] < 1.0:\n",
        "         raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n",
        "        defaults = dict(lr=lr, betas=betas, eps=eps,\n",
        "                     weight_decay=weight_decay, amsgrad=amsgrad)\n",
        "        super(ExtraAdam, self).__init__(params, defaults)\n",
        "\n",
        "    def __setstate__(self, state):\n",
        "        super(ExtraAdam, self).__setstate__(state)\n",
        "        for group in self.param_groups:\n",
        "            group.setdefault('amsgrad', False)\n",
        "\n",
        "    def update(self, p, group):\n",
        "        if p.grad is None:\n",
        "            return None\n",
        "        grad = p.grad.data\n",
        "        if grad.is_sparse:\n",
        "            raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')\n",
        "        amsgrad = group['amsgrad']\n",
        "\n",
        "        state = self.state[p]\n",
        "\n",
        "        # State initialization\n",
        "        if len(state) == 0:\n",
        "            state['step'] = 0\n",
        "            # Exponential moving average of gradient values\n",
        "            state['exp_avg'] = torch.zeros_like(p.data)\n",
        "            # Exponential moving average of squared gradient values\n",
        "            state['exp_avg_sq'] = torch.zeros_like(p.data)\n",
        "            if amsgrad:\n",
        "                # Maintains max of all exp. moving avg. of sq. grad. values\n",
        "                state['max_exp_avg_sq'] = torch.zeros_like(p.data)\n",
        "\n",
        "        exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n",
        "        if amsgrad:\n",
        "            max_exp_avg_sq = state['max_exp_avg_sq']\n",
        "        beta1, beta2 = group['betas']\n",
        "\n",
        "        state['step'] += 1\n",
        "\n",
        "        if group['weight_decay'] != 0:\n",
        "            grad = grad.add(group['weight_decay'], p.data)\n",
        "\n",
        "        # Decay the first and second moment running average coefficient\n",
        "        exp_avg.mul_(beta1).add_(1 - beta1, grad)\n",
        "        exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n",
        "        if amsgrad:\n",
        "            # Maintains the maximum of all 2nd moment running avg. till now\n",
        "            torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)\n",
        "            # Use the max. for normalizing running avg. of gradient\n",
        "            denom = max_exp_avg_sq.sqrt().add_(group['eps'])\n",
        "        else:\n",
        "            denom = exp_avg_sq.sqrt().add_(group['eps'])\n",
        "\n",
        "        bias_correction1 = 1 - beta1 ** state['step']\n",
        "        bias_correction2 = 1 - beta2 ** state['step']\n",
        "        step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1\n",
        "\n",
        "        return -step_size*exp_avg/denom"
      ],
      "metadata": {
        "id": "3JRKv-C9Q2Tq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5RBGK2Jp8eRY"
      },
      "outputs": [],
      "source": [
        "#function for plottings using kernel density estimation. Same as the original implementation.\n",
        "def kde(mu, tau, bbox=[-1.6, 1.6, -1.6, 1.6], save_file=\"\", xlabel=\"\", ylabel=\"\", cmap='Blues'):\n",
        "    values = np.vstack([mu, tau])\n",
        "    kernel = sp.stats.gaussian_kde(values)\n",
        "\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.axis(bbox)\n",
        "    ax.set_aspect(abs(bbox[1]-bbox[0])/abs(bbox[3]-bbox[2]))\n",
        "    ax.set_xlabel(xlabel)\n",
        "    ax.set_ylabel(ylabel)\n",
        "    plt.tick_params(\n",
        "        axis='x',          # changes apply to the x-axis\n",
        "        which='both',      # both major and minor ticks are affected\n",
        "        bottom='off',      # ticks along the bottom edge are off\n",
        "        top='off',         # ticks along the top edge are off\n",
        "        labelbottom='off') # labels along the bottom edge are off\n",
        "    plt.tick_params(\n",
        "        axis='y',          # changes apply to the x-axis\n",
        "        which='both',      # both major and minor ticks are affected\n",
        "        left='off',      # ticks along the bottom edge are off\n",
        "        right='off',         # ticks along the top edge are off\n",
        "        labelleft='off') # labels along the bottom edge are off\n",
        "    \n",
        "    xx, yy = np.mgrid[bbox[0]:bbox[1]:300j, bbox[2]:bbox[3]:300j]\n",
        "    positions = np.vstack([xx.ravel(), yy.ravel()])\n",
        "    f = np.reshape(kernel(positions).T, xx.shape)\n",
        "    cfset = ax.contourf(xx, yy, f, cmap=cmap)\n",
        "\n",
        "    if save_file != \"\":\n",
        "        plt.savefig(save_file, bbox_inches='tight')\n",
        "        plt.close(fig)\n",
        "    else:\n",
        "        plt.show()\n",
        "\n",
        "# function to plot eignevalues in the complex plane. Same as the original implementation.\n",
        "def complex_scatter(points, bbox=None, save_file=\"\", xlabel=\"real part\", ylabel=\"imaginary part\", cmap='Blues'):\n",
        "    fig, ax = plt.subplots()\n",
        "\n",
        "    if bbox is not None:\n",
        "        ax.axis(bbox)\n",
        "\n",
        "    ax.set_xlabel(xlabel)\n",
        "    ax.set_ylabel(ylabel)\n",
        "\n",
        "    xx = [p.real for p in points]\n",
        "    yy = [p.imag for p in points]\n",
        "    \n",
        "    plt.plot(xx, yy, 'X')\n",
        "    plt.grid()\n",
        "\n",
        "    if save_file != \"\":\n",
        "        plt.savefig(save_file, bbox_inches='tight')\n",
        "        plt.close(fig)\n",
        "    else:\n",
        "        plt.show()        "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BeW0m9Y27Caw"
      },
      "outputs": [],
      "source": [
        "from os.path import join\n",
        "def plot_eigens(iteration):\n",
        "    gen_out, real_in, fake_d_out_gen, fake_d_out_disc, fake_d_out, real_d_out = batch_net_outputs()\n",
        "    gen_loss_detached, disc_loss_detached, gen_loss, disc_loss = net_losses(fake_d_out_gen, fake_d_out_disc, fake_d_out, real_d_out)\n",
        "    p_count = torch.cat([x.flatten() for x in params]).shape[0]\n",
        "\n",
        "    gen_net.zero_grad()\n",
        "    gen_grad = autograd.grad(gen_loss, gen_net.parameters(), retain_graph=True, create_graph=True)\n",
        "    disc_net.zero_grad()\n",
        "    disc_grad = autograd.grad(disc_loss, disc_net.parameters(), retain_graph=True, create_graph=True)\n",
        "\n",
        "    v = list(gen_grad) + list(disc_grad)\n",
        "    v = torch.cat([t.flatten() for t in v])\n",
        "    jacobian = torch.zeros([p_count * p_count]).to(device)\n",
        "\n",
        "    for i in range(p_count):\n",
        "        jacobian[i*p_count: (i+1)*p_count] = torch.cat([x.flatten() for x in autograd.grad(v[i], params, retain_graph=True)])\n",
        "\n",
        "    jacobian = jacobian.reshape([p_count, -1])\n",
        "\n",
        "    jacobian2 = jacobian - gamma * torch.mm(jacobian.T, jacobian)\n",
        "\n",
        "    eigens = torch.linalg.eigvals(jacobian)\n",
        "    eigens2 = torch.linalg.eigvals(jacobian2)\n",
        "\n",
        "    print(eigens.shape)\n",
        "    breakpoint()\n",
        "\n",
        "\n",
        "    save_path = join(path, 'Eig_v_' + str(iteration) + \"_\" + method + \"_\" + optim_name + \"_\" + str(z_dim) + \"_\" + str(gamma) + \"_\" + net_size + '.png')\n",
        "    cmap='Blues'\n",
        "    complex_scatter(eigens.cpu().detach().numpy(), bbox=[-1.0, 1.0, -0.15, 0.15], save_file=save_path, cmap=cmap)\n",
        "\n",
        "    save_path = join(path, 'Eig_w_' + str(iteration) + \"_\" + method   + \"_\" + optim_name + \"_\" + str(z_dim) + \"_\" + str(gamma) + \"_\" + net_size +  '.png')\n",
        "    complex_scatter(eigens2.cpu().detach().numpy(), bbox=[-1.0, 1.0, -0.15, 0.15], save_file=save_path, cmap=cmap)\n",
        "\n",
        "\n",
        "def plot_kde(iteration, real_input=False):\n",
        "    z = torch.normal(mean=0, std=1, size=[batch_size * 5, z_dim]).to(device)\n",
        "    inp = gen_net(z)\n",
        "\n",
        "    save_path = join(path, 'KDE', method + \"_\" + str(iteration)  + \"_\" + optim_name + \"_\" + str(z_dim) + \"_\" + str(gamma) + \"_\" + net_size + '.png')\n",
        "    cmap='Blues'\n",
        "    \n",
        "    if real_input:\n",
        "        angles = torch.tensor([2*np.pi*k/8 for k in range(batch_size * 5)]).to(device)\n",
        "        mus = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)\n",
        "        inp = mus + torch.normal(mean=0, std=sigma, size=[batch_size * 5, 2]).to(device)\n",
        "\n",
        "        save_path = join(path, 'KDE', 'original.png')\n",
        "        cmap='Reds'\n",
        "\n",
        "    kde(inp[:, 0].cpu().detach().numpy(), inp[:, 1].cpu().detach().numpy(), save_file=save_path, cmap=cmap)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XQYDRnTsuJVf"
      },
      "outputs": [],
      "source": [
        "class Net(nn.Module):\n",
        "    def __init__(self, indim, outdim):\n",
        "        super(Net, self).__init__()\n",
        "        self.fc1 = nn.Linear(indim, 16)\n",
        "        self.fc2 = nn.Linear(16, 16)\n",
        "        self.fc3 = nn.Linear(16, 16)\n",
        "        self.fc4 = nn.Linear(16, 16)\n",
        "        self.fc5 = nn.Linear(16, outdim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.fc1(x))\n",
        "        out = F.relu(self.fc2(out))\n",
        "        out = F.relu(self.fc3(out))\n",
        "        out = F.relu(self.fc4(out))\n",
        "        out = self.fc5(out)\n",
        "\n",
        "        return out\n",
        "\n",
        "class Net_big(nn.Module):\n",
        "    def __init__(self, indim, outdim):\n",
        "        super(Net_big, self).__init__()\n",
        "        self.fc1 = nn.Linear(indim, 16)\n",
        "        self.fc2 = nn.Linear(16, 16)\n",
        "        self.fc3 = nn.Linear(16, 16)\n",
        "        self.fc4 = nn.Linear(16, 16)\n",
        "        self.fc5 = nn.Linear(16, 16)\n",
        "        self.fc6 = nn.Linear(16, 16)\n",
        "        self.fc7 = nn.Linear(16, 16)\n",
        "        self.fc8 = nn.Linear(16, outdim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.fc1(x))\n",
        "        out = F.relu(self.fc2(out))\n",
        "        out = F.relu(self.fc3(out))\n",
        "        out = F.relu(self.fc4(out))\n",
        "        out = F.relu(self.fc5(out))\n",
        "        out = F.relu(self.fc6(out))\n",
        "        out = F.relu(self.fc7(out))\n",
        "        out = self.fc8(out)\n",
        "\n",
        "        return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5iupuM-WwmLe"
      },
      "outputs": [],
      "source": [
        "criterion = nn.BCEWithLogitsLoss()\n",
        "def net_losses(fake_d_out_gen, fake_d_out_disc, fake_d_out, real_d_out):\n",
        "    disc_loss_detached = criterion(input=fake_d_out_disc, target=torch.zeros_like(fake_d_out_disc))\n",
        "    disc_loss_detached += criterion(real_d_out, target=torch.ones_like(real_d_out))\n",
        "\n",
        "    gen_loss_detached =  criterion(fake_d_out_gen, target=torch.ones_like(fake_d_out_gen))\n",
        "\n",
        "    disc_loss = criterion(input=fake_d_out, target=torch.zeros_like(fake_d_out))\n",
        "    disc_loss += criterion(real_d_out, target=torch.ones_like(real_d_out))\n",
        "\n",
        "    gen_loss =  criterion(fake_d_out, target=torch.ones_like(fake_d_out))\n",
        "    \n",
        "    return gen_loss_detached, disc_loss_detached, gen_loss, disc_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BEXPDP9B4ZPh"
      },
      "outputs": [],
      "source": [
        "#Parameters\n",
        "z_dim = 32\n",
        "gamma = 10.0\n",
        "lr = 1e-4\n",
        "sigma = 1e-2\n",
        "steps = 20000\n",
        "batch_size = 512\n",
        "method = 'ConsOpt' #'SimGA' #'ConsOpt'\n",
        "optim_name = 'Adam'\n",
        "net_size = 'small'\n",
        "\n",
        "lr_adam = 3e-4\n",
        "beta = 0.55\n",
        "alpha = 0.6"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VQcudq8LCSQ-"
      },
      "outputs": [],
      "source": [
        "if net_size == \"small\":\n",
        "  gen_net = Net(z_dim, 2).to(device)\n",
        "  disc_net = Net(2, 1).to(device)\n",
        "else:\n",
        "  gen_net = Net_big(z_dim, 2).to(device)\n",
        "  disc_net = Net_big(2, 1).to(device)\n",
        "\n",
        "params = list(gen_net.parameters()) + list(disc_net.parameters())\n",
        "\n",
        "\n",
        "if optim_name == \"Adam\":\n",
        "  gen_opt = optim.Adam(gen_net.parameters(), lr=lr_adam, betas=(0.5, 0.9))\n",
        "  disc_opt = optim.Adam(disc_net.parameters(), lr=lr_adam, betas=(beta, 0.9))\n",
        "elif optim_name == 'ExtraAdam':\n",
        "  dis_optimizer = ExtraAdam(disc_net.parameters(), lr=2e-4, betas=(0.5, 0.9))\n",
        "  gen_optimizer = ExtraAdam(gen_net.parameters(), lr=2e-5, betas=(0.5, 0.9))\n",
        "else:\n",
        "  gen_opt = optim.RMSprop(gen_net.parameters(), lr=lr)\n",
        "  disc_opt = optim.RMSprop(disc_net.parameters(), lr=lr)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bBSm4fvv7rnV"
      },
      "outputs": [],
      "source": [
        "def batch_net_outputs():\n",
        "    z = torch.normal(mean=0, std=1, size=[batch_size, z_dim]).to(device)\n",
        "    gen_out = gen_net(z)\n",
        "    fake_d_out_gen = copy.deepcopy(disc_net)(gen_out)\n",
        "    fake_d_out_disc = disc_net(gen_out.detach())\n",
        "    fake_d_out = disc_net(gen_out)\n",
        "\n",
        "    angles = torch.tensor([2*np.pi*k/8 for k in range(batch_size)]).to(device)\n",
        "    mus = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)\n",
        "    real_in = mus + torch.normal(mean=0, std=sigma, size=[batch_size, 2]).to(device)\n",
        "    real_d_out = disc_net(real_in)\n",
        "\n",
        "    return gen_out, real_in, fake_d_out_gen, fake_d_out_disc, fake_d_out, real_d_out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zXqSfAm5wQGU"
      },
      "outputs": [],
      "source": [
        "# Uncomment if you needed to load a model\n",
        "\n",
        "# gen_path = join(path, 'Models', 'gen_SimGA_0.pt')\n",
        "# disc_path = join(path, 'Models', 'disc_SimGA_0.pt')\n",
        "# gen_net.load_state_dict(torch.load(gen_path))\n",
        "# disc_net.load_state_dict(torch.load(disc_path))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HJIyfV7h4xk9"
      },
      "outputs": [],
      "source": [
        "# plot_kde(0, real_input=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "f7fb62f653484658bc740e3e597d22c9",
            "74801c66de4b4319b6600d3c32daf523",
            "48d95a94d8724fd9988ed8e1ec755d1b",
            "8503498c9bd14729ba81084a50df97d1",
            "badff3bac8ed456cadd952c09681d376",
            "da11e1fe45a347528251737235ba42a3",
            "ae62a8ccd0ac462b89f8498c986b40cd",
            "c480dab59e3949fe81fbe6a2f8d5b2dc",
            "d039ae3ff53c40b2b49deff35f80d23a",
            "9ea4fcb1afa1450f96f8a4135f028a08",
            "38a891fbf6db45d9a1dd82cff317f6b7"
          ]
        },
        "id": "wlXrjB1y49zX",
        "outputId": "b49a268c-21fa-46af-9f95-bfda4e1aea63"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/20001 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f7fb62f653484658bc740e3e597d22c9"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "PYDEV DEBUGGER WARNING:\n",
            "sys.settrace() should not be used when the debugger is being used.\n",
            "This may cause the debugger to stop working correctly.\n",
            "If this is needed, please check: \n",
            "http://pydev.blogspot.com/2007/06/why-cant-pydev-debugger-work-with.html\n",
            "to see how to restore the debug tracing back correctly.\n",
            "Call Location:\n",
            "  File \"/usr/lib/python3.9/bdb.py\", line 334, in set_trace\n",
            "    sys.settrace(self.trace_dispatch)\n",
            "\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([2259])\n",
            "> <ipython-input-18-d0ca30d8166c>(30)plot_eigens()\n",
            "-> save_path = join(path, 'Eig_v_' + str(iteration) + \"_\" + method + \"_\" + optim_name + \"_\" + str(z_dim) + \"_\" + str(gamma) + \"_\" + net_size + '.png')\n",
            "(Pdb) eigens.shape\n",
            "torch.Size([2259])\n",
            "(Pdb) n\n",
            "> <ipython-input-18-d0ca30d8166c>(31)plot_eigens()\n",
            "-> cmap='Blues'\n",
            "(Pdb) eigen2.shape\n",
            "*** NameError: name 'eigen2' is not defined\n",
            "(Pdb) eigens2\n",
            "tensor([-20.4477+0.0000j,  -1.8414+0.4181j,  -1.8414-0.4181j,  ...,\n",
            "          0.0000+0.0000j,   0.0000+0.0000j,   0.0000+0.0000j], device='cuda:0')\n",
            "(Pdb) eigens2.shape\n",
            "torch.Size([2259])\n",
            "--KeyboardInterrupt--\n",
            "(Pdb) exit\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "PYDEV DEBUGGER WARNING:\n",
            "sys.settrace() should not be used when the debugger is being used.\n",
            "This may cause the debugger to stop working correctly.\n",
            "If this is needed, please check: \n",
            "http://pydev.blogspot.com/2007/06/why-cant-pydev-debugger-work-with.html\n",
            "to see how to restore the debug tracing back correctly.\n",
            "Call Location:\n",
            "  File \"/usr/lib/python3.9/bdb.py\", line 359, in set_quit\n",
            "    sys.settrace(None)\n",
            "\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "BdbQuit",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mBdbQuit\u001b[0m                                   Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-19-b4da37419d9f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m%\u001b[0m\u001b[0;36m5000\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m         \u001b[0;31m# if method == 'ConsOpt':\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m         \u001b[0mplot_eigens\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\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     10\u001b[0m         \u001b[0mplot_kde\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-18-d0ca30d8166c>\u001b[0m in \u001b[0;36mplot_eigens\u001b[0;34m(iteration)\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m     \u001b[0msave_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'Eig_v_'\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miteration\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0moptim_name\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgamma\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mnet_size\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'.png'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m     \u001b[0mcmap\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'Blues'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     32\u001b[0m     \u001b[0mcomplex_scatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0meigens\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\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[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbbox\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1.0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1.0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m0.15\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0.15\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msave_file\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msave_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcmap\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcmap\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-18-d0ca30d8166c>\u001b[0m in \u001b[0;36mplot_eigens\u001b[0;34m(iteration)\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m     \u001b[0msave_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'Eig_v_'\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miteration\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0moptim_name\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgamma\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mnet_size\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'.png'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m     \u001b[0mcmap\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'Blues'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     32\u001b[0m     \u001b[0mcomplex_scatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0meigens\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\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[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbbox\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1.0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1.0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m0.15\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0.15\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msave_file\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msave_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcmap\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcmap\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/lib/python3.9/bdb.py\u001b[0m in \u001b[0;36mtrace_dispatch\u001b[0;34m(self, frame, event, arg)\u001b[0m\n\u001b[1;32m     86\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;31m# None\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'line'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\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     89\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'call'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/lib/python3.9/bdb.py\u001b[0m in \u001b[0;36mdispatch_line\u001b[0;34m(self, frame)\u001b[0m\n\u001b[1;32m    111\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstop_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbreak_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\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    112\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 113\u001b[0;31m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mquitting\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mBdbQuit\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    114\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_dispatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mBdbQuit\u001b[0m: "
          ]
        }
      ],
      "source": [
        "n_gen_update = 0\n",
        "for i in tqdm.notebook.tqdm(range(steps+1)):\n",
        "\n",
        "    gen_out, real_in, fake_d_out_gen, fake_d_out_disc, fake_d_out, real_d_out = batch_net_outputs()\n",
        "    gen_loss_detached, disc_loss_detached, gen_loss, disc_loss = net_losses(fake_d_out_gen, fake_d_out_disc, fake_d_out, real_d_out)\n",
        "\n",
        "    if i%5000 == 0:\n",
        "        # if method == 'ConsOpt':\n",
        "        plot_eigens(i)\n",
        "        plot_kde(i)\n",
        "        \n",
        "        gen_path = join(path, 'Models', 'gen_' + method + \"_\" + str(i) + \"_\" + optim_name + \"_\" + str(z_dim) + \"_\" + str(gamma) + \"_\" + net_size + '.pt')\n",
        "        disc_path = join(path, 'Models', 'disc_'+ method + \"_\" + str(i)  + \"_\" + optim_name + \"_\" + str(z_dim) + \"_\" + str(gamma) + \"_\" + net_size + '.pt')\n",
        "        torch.save(gen_net.state_dict(), gen_path)\n",
        "        torch.save(disc_net.state_dict(), disc_path)      \n",
        "\n",
        "    if method == 'ConsOpt':\n",
        "\n",
        "        gen_net.zero_grad()\n",
        "        gen_grad = autograd.grad(gen_loss, gen_net.parameters(), retain_graph=True, create_graph=True)\n",
        "        disc_net.zero_grad()\n",
        "        disc_grad = autograd.grad(disc_loss, disc_net.parameters(), retain_graph=True, create_graph=True)\n",
        "\n",
        "        v = list(gen_grad) + list(disc_grad)\n",
        "        v = torch.cat([t.flatten() for t in v])\n",
        "\n",
        "        L = 1/2 * torch.dot(v, v)\n",
        "        jgrads = autograd.grad(L, params, retain_graph=True)\n",
        "        \n",
        "        gen_opt.zero_grad()\n",
        "\n",
        "        for i in range(len(params)):\n",
        "               params[i].grad = jgrads[i] * gamma\n",
        "        gen_loss_detached.backward(retain_graph=True, create_graph=True)\n",
        "        gen_opt.step()\n",
        "\n",
        "        disc_opt.zero_grad()\n",
        "\n",
        "        for i in range(len(params)):\n",
        "               params[i].grad = jgrads[i] * gamma\n",
        "        disc_loss_detached.backward(retain_graph=True, create_graph=True)\n",
        "        disc_opt.step()\n",
        "    elif method == \"ExtraAdam\":\n",
        "\n",
        "        for p in gen_net.parameters():\n",
        "            p.requires_grad = False\n",
        "        dis_optimizer.zero_grad()\n",
        "        disc_loss.backward(retain_graph=True)\n",
        "\n",
        "        if (i+1)%2 != 0:\n",
        "            dis_optimizer.extrapolation()\n",
        "        else:\n",
        "            dis_optimizer.step()\n",
        "\n",
        "        for p in gen_net.parameters():\n",
        "            p.requires_grad = True\n",
        "\n",
        "        for p in disc_net.parameters():\n",
        "            p.requires_grad = False\n",
        "        gen_optimizer.zero_grad()\n",
        "        gen_loss.backward()\n",
        "\n",
        "        if (i+1)%2 != 0:\n",
        "            gen_optimizer.extrapolation()\n",
        "        else:\n",
        "            n_gen_update += 1\n",
        "            gen_optimizer.step()\n",
        "      \n",
        "        for p in disc_net.parameters():\n",
        "            p.requires_grad = True\n",
        "    else:\n",
        "        gen_opt.zero_grad()\n",
        "        gen_loss_detached.backward(retain_graph=True, create_graph=True)\n",
        "        gen_opt.step()\n",
        "\n",
        "        disc_opt.zero_grad()\n",
        "        disc_loss_detached.backward(retain_graph=True, create_graph=True)\n",
        "        disc_opt.step()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m97XALSF3lsk",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 222
        },
        "outputId": "5d184a41-93bf-4c53-8c69-baf145fee294"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-16-eb1a18fac8c3>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0meigens\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m: name 'eigens' is not defined"
          ]
        }
      ],
      "source": [
        "eigens"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "5unKlbzrsrco"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "UghmB3pbssVh"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "f7fb62f653484658bc740e3e597d22c9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "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_74801c66de4b4319b6600d3c32daf523",
              "IPY_MODEL_48d95a94d8724fd9988ed8e1ec755d1b",
              "IPY_MODEL_8503498c9bd14729ba81084a50df97d1"
            ],
            "layout": "IPY_MODEL_badff3bac8ed456cadd952c09681d376"
          }
        },
        "74801c66de4b4319b6600d3c32daf523": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "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_da11e1fe45a347528251737235ba42a3",
            "placeholder": "​",
            "style": "IPY_MODEL_ae62a8ccd0ac462b89f8498c986b40cd",
            "value": "  0%"
          }
        },
        "48d95a94d8724fd9988ed8e1ec755d1b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "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": "danger",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c480dab59e3949fe81fbe6a2f8d5b2dc",
            "max": 20001,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d039ae3ff53c40b2b49deff35f80d23a",
            "value": 0
          }
        },
        "8503498c9bd14729ba81084a50df97d1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "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_9ea4fcb1afa1450f96f8a4135f028a08",
            "placeholder": "​",
            "style": "IPY_MODEL_38a891fbf6db45d9a1dd82cff317f6b7",
            "value": " 0/20001 [01:11&lt;?, ?it/s]"
          }
        },
        "badff3bac8ed456cadd952c09681d376": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "da11e1fe45a347528251737235ba42a3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "ae62a8ccd0ac462b89f8498c986b40cd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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": ""
          }
        },
        "c480dab59e3949fe81fbe6a2f8d5b2dc": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "d039ae3ff53c40b2b49deff35f80d23a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "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": ""
          }
        },
        "9ea4fcb1afa1450f96f8a4135f028a08": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "38a891fbf6db45d9a1dd82cff317f6b7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}