notebooks/minimal_code.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/YannDubs/lossyless/blob/main/notebooks/minimal_code.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vDycnS3n8zea"
      },
      "source": [
        "# Minimal Lossyless Code\n",
        "\n",
        "This notebook contains a minimal pipeline for training and evaluating the proposed CLIP compressor in our paper [**Lossy Compression for Lossless Prediction**](https://arxiv.org/pdf/2106.10800.pdf). The focus is on simplicity and understandibility.\n",
        "\n",
        "**Make sure that you use a GPU** (on COLAB: runtime -> change runtime type -> Hardware accelerator: GPU)\n",
        "\n",
        "## Environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lGArtm819IRl",
        "outputId": "4c69c555-2bf4-4060-a620-1391f9f6f357"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "  Building wheel for clip (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "!pip install git+https://github.com/openai/CLIP.git --quiet # pretrained CLIP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vpMI2dWV9tHP",
        "outputId": "21fd329e-5b61-4b00-d4c9-02f18d1afa46"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[33mWARNING: Skipping torchtext as it is not installed.\u001b[0m\n"
          ]
        }
      ],
      "source": [
        "!pip uninstall -y torchtext --quiet # in case it got installed (wouldn't work due to torch version issue)\n",
        "!pip install scikit-learn==0.24.2 lightning-bolts==0.3.4 compressai==1.1.5 pytorch-lightning==1.3.8 --quiet # for evaluation, training, and compression"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pndo0iM59BGZ"
      },
      "source": [
        "Basic variables depending on whether using GPU or not."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "xlRwxV18-Sg7"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "data_dir = \"data/\"\n",
        "if torch.cuda.is_available():\n",
        "    device, precision, gpus = \"cuda\", 16, 1\n",
        "else:\n",
        "    device, precision, gpus = \"cpu\", 32, 0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YwEuObYZ-itL"
      },
      "source": [
        "## Downloads\n",
        "Downloads the necessary data and the pretrained CLIP."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Bv8x6Uey-fj1",
        "outputId": "d67d0ac1-1e48-4ab6-d86d-b40a19eac2eb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n",
            "Files already downloaded and verified\n",
            "Files already downloaded and verified\n"
          ]
        }
      ],
      "source": [
        "import clip\n",
        "from torchvision.datasets import CIFAR10, STL10\n",
        "\n",
        "# pretrained CLIP\n",
        "pretrained, preprocess = clip.load(\"ViT-B/32\", device)\n",
        "\n",
        "# train data\n",
        "cifar = CIFAR10(data_dir, download=True, train=True, transform=preprocess)\n",
        "\n",
        "# eval data\n",
        "stl10_train = STL10(data_dir, download=True, split=\"train\", transform=preprocess)\n",
        "stl10_test = STL10(data_dir, download=True, split=\"test\", transform=preprocess)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HU_D6x2r_Z08"
      },
      "source": [
        "## Training the compressor\n",
        "\n",
        "We will now train the compressor on CIFAR10, by:\n",
        "\n",
        "1. Featurize the training dataset by the pretrained CLIP => instead of working with raw images we directly work with features to be mcuh quicker.\n",
        "2. Create our entropy bottleneck model.\n",
        "3. Train the entropy bottleneck on CIFAR10.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "HJnSSrDY_Dvi"
      },
      "outputs": [],
      "source": [
        "import tqdm\n",
        "from torch.utils.data import DataLoader\n",
        "import numpy as np\n",
        "\n",
        "def clip_featurize_data(dataset, device, pretrained):\n",
        "    \"\"\"Featurize a dataset using the pretrained CLIP model.\"\"\"\n",
        "    with torch.no_grad():\n",
        "        Z, Y = [], []\n",
        "        for x, y in tqdm.tqdm(DataLoader(dataset, batch_size=128, num_workers=16)):\n",
        "            Z += [pretrained.encode_image(x.to(device).half()).cpu().numpy()]\n",
        "            Y += [y.cpu().numpy()]\n",
        "    return np.concatenate(Z), np.concatenate(Y)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "5853hnkXAvoV"
      },
      "outputs": [],
      "source": [
        "from torch.optim import Adam\n",
        "from torch.optim.lr_scheduler import StepLR\n",
        "import pytorch_lightning as pl\n",
        "from compressai.entropy_models import EntropyBottleneck\n",
        "import math\n",
        "\n",
        "class ArrayCompressor(pl.LightningModule):\n",
        "    \"\"\"Compressor for any vectors, by using an entropy bottleneck and MSE distortion.\"\"\"\n",
        "\n",
        "    def __init__(self, *args, **kwargs):\n",
        "        super().__init__()\n",
        "        self.save_hyperparameters()\n",
        "        self.bottleneck = EntropyBottleneck(self.hparams.z_dim)\n",
        "        self.scaling = torch.nn.Parameter(torch.ones(self.hparams.z_dim))\n",
        "        self.biasing = torch.nn.Parameter(torch.zeros(self.hparams.z_dim))\n",
        "        self.is_updated = False\n",
        "\n",
        "    def forward(self, batch):\n",
        "        z, y = batch\n",
        "        z = (z + self.biasing) * self.scaling.exp()\n",
        "        z_hat, q_z = self.bottleneck(z.unsqueeze(-1).unsqueeze(-1))\n",
        "        z_hat = z_hat.squeeze() / self.scaling.exp() - self.biasing\n",
        "        return z_hat, q_z.squeeze(), y.squeeze()\n",
        "\n",
        "    def step(self, batch, *args, **kwargs):\n",
        "        z_hat, q_z, _ = self(batch)\n",
        "        rate = -torch.log(q_z).sum(-1).mean()\n",
        "        distortion = torch.norm(batch[0] - z_hat, p=1, dim=-1).mean()\n",
        "        self.log_dict(\n",
        "            {\"rate\": rate / math.log(2), \"distortion\": distortion}, prog_bar=True\n",
        "        )\n",
        "        return distortion + self.hparams.lmbda * rate\n",
        "\n",
        "    def training_step(self, batch, _, optimizer_idx=0):\n",
        "        return self.step(batch) if optimizer_idx == 0 else self.bottleneck.loss()\n",
        "\n",
        "    def predict_step(self, batch, _, __):\n",
        "        return self.compress(batch[0]), batch[1].cpu().numpy()\n",
        "\n",
        "    def compress(self, z):\n",
        "        if not self.is_updated:\n",
        "            self.bottleneck.update(force=True)\n",
        "            self.is_updated = True\n",
        "        z = (z + self.biasing) * self.scaling.exp()\n",
        "        return self.bottleneck.compress(z.unsqueeze(-1).unsqueeze(-1))\n",
        "\n",
        "    def decompress(self, z_bytes):\n",
        "        z_hat = self.bottleneck.decompress(z_bytes, [1, 1]).squeeze()\n",
        "        return (z_hat / self.scaling.exp()) - self.biasing\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        param = [p for n, p in self.named_parameters() if not n.endswith(\".quantiles\")]\n",
        "        quantile_param = [\n",
        "            p for n, p in self.named_parameters() if n.endswith(\".quantiles\")\n",
        "        ]\n",
        "        optimizer = Adam(param, lr=self.hparams.lr)\n",
        "        optimizer_coder = Adam(quantile_param, lr=self.hparams.lr)\n",
        "        scheduler = StepLR(optimizer, self.hparams.lr_step)\n",
        "        scheduler_coder = StepLR(optimizer_coder, self.hparams.lr_step)\n",
        "        return [optimizer, optimizer_coder], [scheduler, scheduler_coder]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 355,
          "referenced_widgets": [
            "a4c3872442ba4c97a1904b7c358b1a10",
            "95b9ec5768b740629c711e57e74e6eea",
            "b2ad31cb2b234588acd7eeef639be546",
            "3d387cdc6a574cbc991fc8fd9b78cbd8",
            "605a38d924454f0b9fcb9181bf8dd19d",
            "d82e2b9fccff42e3b08ee07471bcb555",
            "31f6cc44c5b4415bb6fe2adbb6037714",
            "cc57814858bb4525b5426558e1965200"
          ]
        },
        "id": "ar-3pEaAJRqM",
        "outputId": "6b102820-0125-4842-a959-886dfc891ede"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 391/391 [01:47<00:00,  3.65it/s]\n",
            "GPU available: True, used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "Using native 16bit precision.\n",
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/configuration_validator.py:99: UserWarning: you passed in a val_dataloader but have no validation_step. Skipping val loop\n",
            "  rank_zero_warn(f'you passed in a {loader_name} but have no {step_name}. Skipping {stage} loop')\n",
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
            "\n",
            "  | Name       | Type              | Params\n",
            "-------------------------------------------------\n",
            "0 | bottleneck | EntropyBottleneck | 31.2 K\n",
            "-------------------------------------------------\n",
            "32.3 K    Trainable params\n",
            "0         Non-trainable params\n",
            "32.3 K    Total params\n",
            "0.129     Total estimated model params size (MB)\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a4c3872442ba4c97a1904b7c358b1a10",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Compressor trained in 3 minutes.\n"
          ]
        }
      ],
      "source": [
        "import time\n",
        "from pl_bolts.datamodules import SklearnDataModule\n",
        "\n",
        "\n",
        "start = time.time()\n",
        "Z_cifar, Y_cifar = clip_featurize_data(cifar, device, pretrained)\n",
        "data_kwargs = dict(\n",
        "    num_workers=16, batch_size=128, pin_memory=True, val_split=0.0, test_split=0\n",
        ")\n",
        "dm_cifar = SklearnDataModule(Z_cifar, Y_cifar, **data_kwargs)\n",
        "compressor = ArrayCompressor(z_dim=512, lmbda=4e-2, lr=1e-1, lr_step=2)\n",
        "trainer = pl.Trainer(gpus=gpus, precision=precision, max_epochs=10, logger=False)\n",
        "trainer.fit(compressor, datamodule=dm_cifar)\n",
        "print(f\"Compressor trained in {(time.time() - start)/60:.0f} minutes.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t5vyp9IGJ5S2"
      },
      "source": [
        "## Evaluating the compressor\n",
        "\n",
        "Now we will evaluate the compressor by quantifying how well it can compress the STL10 dataset and how well a linear classifier can predict from the compressed dataset, by:\n",
        "\n",
        "1. Entropy coding (i.e., compressing) the evaluating data.\n",
        "2. Decompressing the evaluation data.\n",
        "3. Fitting a linear SVM on the training set of the compressed STL10.\n",
        "4. Evaluating the classifier on the test set of the compressed STL10.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "cHJfjbIBLiyV"
      },
      "outputs": [],
      "source": [
        "def compress_data(trainer, dataset, device, pretrained, **kwargs):\n",
        "    \"\"\"Compresses the data using an entropy coder.\"\"\"\n",
        "    start = time.time()\n",
        "    Z, Y = clip_featurize_data(dataset, device, pretrained)\n",
        "    dm = SklearnDataModule(Z, Y, **kwargs)\n",
        "    out = trainer.predict(dataloaders=dm.train_dataloader())\n",
        "    Z_bytes = [o[0] for o in out]\n",
        "    flat_z = [i for batch in Z_bytes for i in batch]\n",
        "    Y = np.concatenate([o[1] for o in out], axis=0)\n",
        "    coding_rate = sum([len(s) for s in flat_z]) * 8 / len(flat_z)\n",
        "    sec_per_img = (time.time() - start) / len(flat_z)\n",
        "    return Z_bytes, Y, coding_rate, sec_per_img\n",
        "\n",
        "\n",
        "def decompress_data(compressor, Z_bytes):\n",
        "    \"\"\"Compresses the data that was entropy coded.\"\"\"\n",
        "    start = time.time()\n",
        "    with torch.no_grad():\n",
        "        Z_hat = [compressor.decompress(b).cpu().numpy() for b in Z_bytes]\n",
        "    sec_per_img = (time.time() - start) / len(Z_hat)\n",
        "    return np.concatenate(Z_hat), sec_per_img"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 234,
          "referenced_widgets": [
            "b68767c3a2964426a2484a1ff3c06087",
            "ecd1af7b6b384171ace76058eb1183b0",
            "e1f168c8717045c6b00a99c046bb082c",
            "aa1f6f40c497486aa6dc512209f952b9",
            "915302eb28aa4663b06e5ca050ad7077",
            "28a92c419ae149c2bcbd6b16a107b0ea",
            "76831be1907a474ebf511a077da69f65",
            "8e78e79d80914ee6a09a6c07de6e995c",
            "19ff31a4288e4cc89e961c45805a0bf6",
            "f89741d2e0b7445b9e3d25304d23debe",
            "c91d188431bf4573b8542295bb50260a",
            "ebabfbfe8d7c44edb56641b9355ea7dc",
            "19450b6f1ee644dcbe47d8bdcf5429c6",
            "b8c158d0481f4ac3875786e4de31e617",
            "fe543739b6554a5bb57b50deef36b829",
            "865e56f3a13c456597666b4486ceb8f7"
          ]
        },
        "id": "QQarqt0dBiBz",
        "outputId": "6bcbfc81-f25a-4070-c1a7-baa38e3dfe7f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 40/40 [00:14<00:00,  2.69it/s]\n",
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/data_loading.py:349: UserWarning: Your predict_dataloader has `shuffle=True`, it is best practice to turn this off for val/test/predict dataloaders.\n",
            "  f'Your {mode}_dataloader has `shuffle=True`, it is best practice to turn'\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b68767c3a2964426a2484a1ff3c06087",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Predicting', layout=Layout(flex='2'), m…"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\r  0%|          | 0/63 [00:00<?, ?it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 63/63 [00:21<00:00,  2.93it/s]\n",
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "19ff31a4288e4cc89e961c45805a0bf6",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Predicting', layout=Layout(flex='2'), m…"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Bit-rate: 1705.8. \t Compression: 222.1 img/sec.\n"
          ]
        }
      ],
      "source": [
        "# entropy code evaluation data. Rate: 1703.6 bits, Compression: 230.2 img/sec\n",
        "Z_b_train, Y_train, *_ = compress_data(trainer, stl10_train, device, pretrained, **data_kwargs)\n",
        "Z_b_test, Y_test, rate, enc_time = compress_data(\n",
        "    trainer, stl10_test, device, pretrained, **data_kwargs\n",
        ")\n",
        "print(f\"Bit-rate: {rate:.1f}. \\t Compression: {1/enc_time:.1f} img/sec.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ppz9yEfQLFDE",
        "outputId": "5c3b8913-49bb-4949-eb30-aa29d4d2b71f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Decompression: 5.3 img/sec.\n"
          ]
        }
      ],
      "source": [
        "# Decompress data. Decoding: 2.8 img/sec (no batch processing)\n",
        "\n",
        "Z_train, _ = decompress_data(compressor, Z_b_train)\n",
        "Z_test, dec_time = decompress_data(compressor, Z_b_test)\n",
        "print(f\"Decompression: {1/dec_time:.1f} img/sec.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Mi01eb0QLH3Q",
        "outputId": "32d01087-2770-49b0-a1b1-975256a0838f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downstream STL10 accuracy: 98.62%.  \t Training time: 0.5 \n"
          ]
        }
      ],
      "source": [
        "# Downstream evaluation. Accuracy: 98.65%  \t Training time: 0.5\n",
        "\n",
        "from sklearn.svm import LinearSVC\n",
        "\n",
        "clf = LinearSVC(C=4e-3)\n",
        "start = time.time()\n",
        "clf.fit(Z_train, Y_train)\n",
        "delta_time = time.time() - start\n",
        "acc = clf.score(Z_test, Y_test)\n",
        "print(\n",
        "    f\"Downstream STL10 accuracy: {acc*100:.2f}%.  \\t Training time: {delta_time:.1f} \"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "V7yZdlEtMB9s"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "minimal_code.ipynb",
      "provenance": []
    },
    "interpreter": {
      "hash": "e50ed95b49794fc89c9e75957291f0e8c9f9bcc9dab75d9e1df0111b44a5a66e"
    },
    "kernelspec": {
      "display_name": "Python 3.8.5 64-bit ('lossyless': conda)",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": ""
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "19450b6f1ee644dcbe47d8bdcf5429c6": {
          "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"
          }
        },
        "19ff31a4288e4cc89e961c45805a0bf6": {
          "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_c91d188431bf4573b8542295bb50260a",
              "IPY_MODEL_ebabfbfe8d7c44edb56641b9355ea7dc"
            ],
            "layout": "IPY_MODEL_f89741d2e0b7445b9e3d25304d23debe"
          }
        },
        "28a92c419ae149c2bcbd6b16a107b0ea": {
          "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": "2",
            "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
          }
        },
        "31f6cc44c5b4415bb6fe2adbb6037714": {
          "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": ""
          }
        },
        "3d387cdc6a574cbc991fc8fd9b78cbd8": {
          "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_cc57814858bb4525b5426558e1965200",
            "placeholder": "​",
            "style": "IPY_MODEL_31f6cc44c5b4415bb6fe2adbb6037714",
            "value": " 391/391 [00:07&lt;00:00, 51.33it/s, loss=39, rate=1.33e+3, distortion=19.50]"
          }
        },
        "605a38d924454f0b9fcb9181bf8dd19d": {
          "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"
          }
        },
        "76831be1907a474ebf511a077da69f65": {
          "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": ""
          }
        },
        "865e56f3a13c456597666b4486ceb8f7": {
          "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
          }
        },
        "8e78e79d80914ee6a09a6c07de6e995c": {
          "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
          }
        },
        "915302eb28aa4663b06e5ca050ad7077": {
          "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"
          }
        },
        "95b9ec5768b740629c711e57e74e6eea": {
          "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": "inline-flex",
            "flex": null,
            "flex_flow": "row wrap",
            "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": "100%"
          }
        },
        "a4c3872442ba4c97a1904b7c358b1a10": {
          "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_b2ad31cb2b234588acd7eeef639be546",
              "IPY_MODEL_3d387cdc6a574cbc991fc8fd9b78cbd8"
            ],
            "layout": "IPY_MODEL_95b9ec5768b740629c711e57e74e6eea"
          }
        },
        "aa1f6f40c497486aa6dc512209f952b9": {
          "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_8e78e79d80914ee6a09a6c07de6e995c",
            "placeholder": "​",
            "style": "IPY_MODEL_76831be1907a474ebf511a077da69f65",
            "value": " 391/? [00:35&lt;00:00, 11.09it/s]"
          }
        },
        "b2ad31cb2b234588acd7eeef639be546": {
          "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": "Epoch 9: 100%",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d82e2b9fccff42e3b08ee07471bcb555",
            "max": 391,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_605a38d924454f0b9fcb9181bf8dd19d",
            "value": 391
          }
        },
        "b68767c3a2964426a2484a1ff3c06087": {
          "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_e1f168c8717045c6b00a99c046bb082c",
              "IPY_MODEL_aa1f6f40c497486aa6dc512209f952b9"
            ],
            "layout": "IPY_MODEL_ecd1af7b6b384171ace76058eb1183b0"
          }
        },
        "b8c158d0481f4ac3875786e4de31e617": {
          "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": "2",
            "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
          }
        },
        "c91d188431bf4573b8542295bb50260a": {
          "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": "Predicting: ",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b8c158d0481f4ac3875786e4de31e617",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_19450b6f1ee644dcbe47d8bdcf5429c6",
            "value": 1
          }
        },
        "cc57814858bb4525b5426558e1965200": {
          "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
          }
        },
        "d82e2b9fccff42e3b08ee07471bcb555": {
          "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": "2",
            "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
          }
        },
        "e1f168c8717045c6b00a99c046bb082c": {
          "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": "Predicting: ",
            "description_tooltip": null,
            "layout": "IPY_MODEL_28a92c419ae149c2bcbd6b16a107b0ea",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_915302eb28aa4663b06e5ca050ad7077",
            "value": 1
          }
        },
        "ebabfbfe8d7c44edb56641b9355ea7dc": {
          "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_865e56f3a13c456597666b4486ceb8f7",
            "placeholder": "​",
            "style": "IPY_MODEL_fe543739b6554a5bb57b50deef36b829",
            "value": " 391/? [00:14&lt;00:00, 27.07it/s]"
          }
        },
        "ecd1af7b6b384171ace76058eb1183b0": {
          "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": "inline-flex",
            "flex": null,
            "flex_flow": "row wrap",
            "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": "100%"
          }
        },
        "f89741d2e0b7445b9e3d25304d23debe": {
          "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": "inline-flex",
            "flex": null,
            "flex_flow": "row wrap",
            "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": "100%"
          }
        },
        "fe543739b6554a5bb57b50deef36b829": {
          "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": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}