reproduce.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "reproduce.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "_lhrn5O-qUYZ",
        "KA1gPmE1h9pI",
        "LUxfDJw1qUYi",
        "2vbPDqd6qUYj",
        "CcEP5cEZqUYl",
        "GSNW-nZCJ4Q0",
        "QRZB9KXyVvfa",
        "MB92Yq_HHy2z",
        "VVZgkJrMbU_x",
        "-xGJF_UDrbDY",
        "anIiMJaJRxeZ",
        "RM4AYreOcG4V",
        "ltVzqPsKtLSW",
        "AU04WKoatRcL"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.12"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_lhrn5O-qUYZ"
      },
      "source": [
        "# Import and misc"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "meO-Mp9jiAFC",
        "outputId": "5bdde34c-9e85-40c0-c119-c76bea0afe16"
      },
      "source": [
        "!pip install torchaudio==0.9.1\n",
        "!pip install thop"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting torchaudio==0.9.1\n",
            "  Downloading torchaudio-0.9.1-cp37-cp37m-manylinux1_x86_64.whl (1.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.9 MB 5.0 MB/s \n",
            "\u001b[?25hCollecting torch==1.9.1\n",
            "  Downloading torch-1.9.1-cp37-cp37m-manylinux1_x86_64.whl (831.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 831.4 MB 6.5 kB/s \n",
            "\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.9.1->torchaudio==0.9.1) (3.10.0.2)\n",
            "Installing collected packages: torch, torchaudio\n",
            "  Attempting uninstall: torch\n",
            "    Found existing installation: torch 1.10.0+cu111\n",
            "    Uninstalling torch-1.10.0+cu111:\n",
            "      Successfully uninstalled torch-1.10.0+cu111\n",
            "  Attempting uninstall: torchaudio\n",
            "    Found existing installation: torchaudio 0.10.0+cu111\n",
            "    Uninstalling torchaudio-0.10.0+cu111:\n",
            "      Successfully uninstalled torchaudio-0.10.0+cu111\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "torchvision 0.11.1+cu111 requires torch==1.10.0, but you have torch 1.9.1 which is incompatible.\n",
            "torchtext 0.11.0 requires torch==1.10.0, but you have torch 1.9.1 which is incompatible.\u001b[0m\n",
            "Successfully installed torch-1.9.1 torchaudio-0.9.1\n",
            "Collecting thop\n",
            "  Downloading thop-0.0.31.post2005241907-py3-none-any.whl (8.7 kB)\n",
            "Requirement already satisfied: torch>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from thop) (1.9.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.0.0->thop) (3.10.0.2)\n",
            "Installing collected packages: thop\n",
            "Successfully installed thop-0.0.31.post2005241907\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bbUpoArCqUYa"
      },
      "source": [
        "from typing import Tuple, Union, List, Callable, Optional\n",
        "from tqdm import tqdm\n",
        "from itertools import islice\n",
        "import pathlib\n",
        "import dataclasses\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch import nn\n",
        "from torch import distributions\n",
        "from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler\n",
        "from torch.nn.utils.rnn import pad_sequence\n",
        "\n",
        "import torchaudio\n",
        "from IPython import display as display_\n",
        "\n",
        "import os"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zONAEKpJFNeu"
      },
      "source": [
        "import random\n",
        "def set_seed(SEED = 777):\n",
        "    torch.manual_seed(SEED)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = False\n",
        "    np.random.seed(SEED)\n",
        "    random.seed(SEED)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "812GwLfqqUYf"
      },
      "source": [
        "# Task"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i1DuQIyRqUYf"
      },
      "source": [
        "In this notebook we will implement a model for finding a keyword in a stream.\n",
        "\n",
        "We will implement the version with CRNN because it is easy and improves the model. \n",
        "(from https://www.dropbox.com/s/22ah2ba7dug6pzw/KWS_Attention.pdf)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8PdhApeEh9pH"
      },
      "source": [
        "@dataclasses.dataclass\n",
        "class TaskConfig:\n",
        "    keyword: str = 'sheila'  # We will use 1 key word -- 'sheila'\n",
        "    batch_size: int = 128\n",
        "    learning_rate: float = 3e-4\n",
        "    weight_decay: float = 1e-5\n",
        "    num_epochs: int = 20\n",
        "    n_mels: int = 40\n",
        "    cnn_out_channels: int = 8\n",
        "    kernel_size: Tuple[int, int] = (5, 20)\n",
        "    stride: Tuple[int, int] = (2, 8)\n",
        "    hidden_size: int = 64\n",
        "    gru_num_layers: int = 2\n",
        "    bidirectional: bool = False\n",
        "    num_classes: int = 2\n",
        "    sample_rate: int = 16000\n",
        "    device: torch.device = torch.device(\n",
        "        'cuda:0' if torch.cuda.is_available() else 'cpu')"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KA1gPmE1h9pI"
      },
      "source": [
        "# Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y2N8zcx9MF1X",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "12502b31-8858-44c6-e351-0b7bfa9bd6a6"
      },
      "source": [
        "!wget http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz -O speech_commands_v0.01.tar.gz\n",
        "!mkdir speech_commands && tar -C speech_commands -xvzf speech_commands_v0.01.tar.gz 1> log"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2021-11-22 21:12:07--  http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz\n",
            "Resolving download.tensorflow.org (download.tensorflow.org)... 172.253.123.128, 2607:f8b0:400c:c16::80\n",
            "Connecting to download.tensorflow.org (download.tensorflow.org)|172.253.123.128|:80... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1489096277 (1.4G) [application/gzip]\n",
            "Saving to: ‘speech_commands_v0.01.tar.gz’\n",
            "\n",
            "speech_commands_v0. 100%[===================>]   1.39G   235MB/s    in 5.9s    \n",
            "\n",
            "2021-11-22 21:12:14 (239 MB/s) - ‘speech_commands_v0.01.tar.gz’ saved [1489096277/1489096277]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "12wBTK0mNUsG"
      },
      "source": [
        "class SpeechCommandDataset(Dataset):\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        transform: Optional[Callable] = None,\n",
        "        path2dir: str = None,\n",
        "        keywords: Union[str, List[str]] = None,\n",
        "        csv: Optional[pd.DataFrame] = None\n",
        "    ):        \n",
        "        self.transform = transform\n",
        "\n",
        "        if csv is None:\n",
        "            path2dir = pathlib.Path(path2dir)\n",
        "            keywords = keywords if isinstance(keywords, list) else [keywords]\n",
        "            \n",
        "            all_keywords = [\n",
        "                p.stem for p in path2dir.glob('*')\n",
        "                if p.is_dir() and not p.stem.startswith('_')\n",
        "            ]\n",
        "\n",
        "            triplets = []\n",
        "            for keyword in all_keywords:\n",
        "                paths = (path2dir / keyword).rglob('*.wav')\n",
        "                if keyword in keywords:\n",
        "                    for path2wav in paths:\n",
        "                        triplets.append((path2wav.as_posix(), keyword, 1))\n",
        "                else:\n",
        "                    for path2wav in paths:\n",
        "                        triplets.append((path2wav.as_posix(), keyword, 0))\n",
        "            \n",
        "            self.csv = pd.DataFrame(\n",
        "                triplets,\n",
        "                columns=['path', 'keyword', 'label']\n",
        "            )\n",
        "\n",
        "        else:\n",
        "            self.csv = csv\n",
        "    \n",
        "    def __getitem__(self, index: int):\n",
        "        instance = self.csv.iloc[index]\n",
        "\n",
        "        path2wav = instance['path']\n",
        "        wav, sr = torchaudio.load(path2wav)\n",
        "        wav = wav.sum(dim=0)\n",
        "        \n",
        "        if self.transform:\n",
        "            wav = self.transform(wav)\n",
        "\n",
        "        return {\n",
        "            'wav': wav,\n",
        "            'keywors': instance['keyword'],\n",
        "            'label': instance['label']\n",
        "        }\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.csv)"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-1rVkT81Pk90"
      },
      "source": [
        "dataset = SpeechCommandDataset(\n",
        "    path2dir='speech_commands', keywords=TaskConfig.keyword\n",
        ")"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "id": "DFwhAXdfQLIA",
        "outputId": "c379b4d8-647f-4140-9d3e-40e6e29c3062"
      },
      "source": [
        "dataset.csv.sample(5)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>path</th>\n",
              "      <th>keyword</th>\n",
              "      <th>label</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>21139</th>\n",
              "      <td>speech_commands/down/afe0b87d_nohash_0.wav</td>\n",
              "      <td>down</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7172</th>\n",
              "      <td>speech_commands/left/3ac2e76f_nohash_0.wav</td>\n",
              "      <td>left</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>45897</th>\n",
              "      <td>speech_commands/nine/d7559b2d_nohash_1.wav</td>\n",
              "      <td>nine</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>49459</th>\n",
              "      <td>speech_commands/go/48bdc11c_nohash_0.wav</td>\n",
              "      <td>go</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1026</th>\n",
              "      <td>speech_commands/sheila/b0c0197e_nohash_0.wav</td>\n",
              "      <td>sheila</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                                               path keyword  label\n",
              "21139    speech_commands/down/afe0b87d_nohash_0.wav    down      0\n",
              "7172     speech_commands/left/3ac2e76f_nohash_0.wav    left      0\n",
              "45897    speech_commands/nine/d7559b2d_nohash_1.wav    nine      0\n",
              "49459      speech_commands/go/48bdc11c_nohash_0.wav      go      0\n",
              "1026   speech_commands/sheila/b0c0197e_nohash_0.wav  sheila      1"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LUxfDJw1qUYi"
      },
      "source": [
        "### Augmentations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dkmkxPWQqUYe"
      },
      "source": [
        "class AugsCreation:\n",
        "\n",
        "    def __init__(self):\n",
        "        self.background_noises = [\n",
        "            'speech_commands/_background_noise_/white_noise.wav',\n",
        "            'speech_commands/_background_noise_/dude_miaowing.wav',\n",
        "            'speech_commands/_background_noise_/doing_the_dishes.wav',\n",
        "            'speech_commands/_background_noise_/exercise_bike.wav',\n",
        "            'speech_commands/_background_noise_/pink_noise.wav',\n",
        "            'speech_commands/_background_noise_/running_tap.wav'\n",
        "        ]\n",
        "\n",
        "        self.noises = [\n",
        "            torchaudio.load(p)[0].squeeze()\n",
        "            for p in self.background_noises\n",
        "        ]\n",
        "\n",
        "    def add_rand_noise(self, audio):\n",
        "\n",
        "        # randomly choose noise\n",
        "        noise_num = torch.randint(low=0, high=len(\n",
        "            self.background_noises), size=(1,)).item()\n",
        "        noise = self.noises[noise_num]\n",
        "\n",
        "        noise_level = torch.Tensor([1])  # [0, 40]\n",
        "\n",
        "        noise_energy = torch.norm(noise)\n",
        "        audio_energy = torch.norm(audio)\n",
        "        alpha = (audio_energy / noise_energy) * \\\n",
        "            torch.pow(10, -noise_level / 20)\n",
        "\n",
        "        start = torch.randint(\n",
        "            low=0,\n",
        "            high=max(int(noise.size(0) - audio.size(0) - 1), 1),\n",
        "            size=(1,)\n",
        "        ).item()\n",
        "        noise_sample = noise[start: start + audio.size(0)]\n",
        "\n",
        "        audio_new = audio + alpha * noise_sample\n",
        "        audio_new.clamp_(-1, 1)\n",
        "        return audio_new\n",
        "\n",
        "    def __call__(self, wav):\n",
        "        aug_num = torch.randint(low=0, high=4, size=(1,)).item()   # choose 1 random aug from augs\n",
        "        augs = [\n",
        "            lambda x: x,\n",
        "            lambda x: (x + distributions.Normal(0, 0.01).sample(x.size())).clamp_(-1, 1),\n",
        "            lambda x: torchaudio.transforms.Vol(.25)(x),\n",
        "            lambda x: self.add_rand_noise(x)\n",
        "        ]\n",
        "\n",
        "        return augs[aug_num](wav)"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ClWThxyYh9pM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0be15d60-22fe-43cd-ec44-74570dbd38c5"
      },
      "source": [
        "set_seed()\n",
        "indexes = torch.randperm(len(dataset))\n",
        "print(indexes[:5])\n",
        "train_indexes = indexes[:int(len(dataset) * 0.8)]\n",
        "val_indexes = indexes[int(len(dataset) * 0.8):]\n",
        "\n",
        "train_df = dataset.csv.iloc[train_indexes].reset_index(drop=True)\n",
        "val_df = dataset.csv.iloc[val_indexes].reset_index(drop=True)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([62005, 57280, 25038,  4203,  4490])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PDPLht5fqUYe"
      },
      "source": [
        "# Sample is a dict of utt, word and label\n",
        "train_set = SpeechCommandDataset(csv=train_df, transform=AugsCreation())\n",
        "val_set = SpeechCommandDataset(csv=val_df)"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mmrJd8WIhkLP"
      },
      "source": [
        ""
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2vbPDqd6qUYj"
      },
      "source": [
        "### Sampler for oversampling:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rfnjRKo2qUYj"
      },
      "source": [
        "# We should provide to WeightedRandomSampler _weight for every sample_; by default it is 1/len(target)\n",
        "\n",
        "def get_sampler(target):\n",
        "    class_sample_count = np.array(\n",
        "        [len(np.where(target == t)[0]) for t in np.unique(target)])   # for every class count it's number of occ.\n",
        "    weight = 1. / class_sample_count\n",
        "    samples_weight = np.array([weight[t] for t in target])\n",
        "    samples_weight = torch.from_numpy(samples_weight)\n",
        "    samples_weigth = samples_weight.float()\n",
        "    sampler = WeightedRandomSampler(samples_weight, len(samples_weight))\n",
        "    return sampler"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UM8gLmHeqUYj"
      },
      "source": [
        "train_sampler = get_sampler(train_set.csv['label'].values)"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lyBqbxp0h9pO"
      },
      "source": [
        "class Collator:\n",
        "    \n",
        "    def __call__(self, data):\n",
        "        wavs = []\n",
        "        labels = []    \n",
        "\n",
        "        for el in data:\n",
        "            wavs.append(el['wav'])\n",
        "            labels.append(el['label'])\n",
        "\n",
        "        # torch.nn.utils.rnn.pad_sequence takes list(Tensors) and returns padded (with 0.0) Tensor\n",
        "        wavs = pad_sequence(wavs, batch_first=True)    \n",
        "        labels = torch.Tensor(labels).long()\n",
        "        return wavs, labels"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e8G9xPRVqUYk"
      },
      "source": [
        "###  Dataloaders"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6wGBMcQiqUYk"
      },
      "source": [
        "# Here we are obliged to use shuffle=False because of our sampler with randomness inside.\n",
        "\n",
        "train_loader = DataLoader(train_set, batch_size=TaskConfig.batch_size,\n",
        "                          shuffle=False, collate_fn=Collator(),\n",
        "                          sampler=train_sampler,\n",
        "                          num_workers=2, pin_memory=True)\n",
        "\n",
        "val_loader = DataLoader(val_set, batch_size=TaskConfig.batch_size,\n",
        "                        shuffle=False, collate_fn=Collator(),\n",
        "                        num_workers=2, pin_memory=True)"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kTlsn6cpqUYk"
      },
      "source": [
        "### Creating MelSpecs on GPU for speeeed: "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pRXMt6it56fW"
      },
      "source": [
        "class LogMelspec:\n",
        "\n",
        "    def __init__(self, is_train, config):\n",
        "        # with augmentations\n",
        "        if is_train:\n",
        "            self.melspec = nn.Sequential(\n",
        "                torchaudio.transforms.MelSpectrogram(\n",
        "                    sample_rate=config.sample_rate,\n",
        "                    n_fft=400,\n",
        "                    win_length=400,\n",
        "                    hop_length=160,\n",
        "                    n_mels=config.n_mels\n",
        "                ),\n",
        "                torchaudio.transforms.FrequencyMasking(freq_mask_param=15),\n",
        "                torchaudio.transforms.TimeMasking(time_mask_param=35),\n",
        "            ).to(config.device)\n",
        "\n",
        "        # no augmentations\n",
        "        else:\n",
        "            self.melspec = torchaudio.transforms.MelSpectrogram(\n",
        "                sample_rate=config.sample_rate,\n",
        "                n_fft=400,\n",
        "                win_length=400,\n",
        "                hop_length=160,\n",
        "                n_mels=config.n_mels\n",
        "            ).to(config.device)\n",
        "\n",
        "    def __call__(self, batch):\n",
        "        # already on device\n",
        "        return torch.log(self.melspec(batch).clamp_(min=1e-9, max=1e9))"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pqkz4_gn8BiF"
      },
      "source": [
        "melspec_train = LogMelspec(is_train=True, config=TaskConfig)\n",
        "melspec_val = LogMelspec(is_train=False, config=TaskConfig)"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zoAxmihY8yxr"
      },
      "source": [
        "### Quality measurment functions:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "euwD1UyuqUYk"
      },
      "source": [
        "# FA - true: 0, model: 1\n",
        "# FR - true: 1, model: 0\n",
        "\n",
        "def count_FA_FR(preds, labels):\n",
        "    FA = torch.sum(preds[labels == 0])\n",
        "    FR = torch.sum(labels[preds == 0])\n",
        "    \n",
        "    # torch.numel - returns total number of elements in tensor\n",
        "    return FA.item() / torch.numel(preds), FR.item() / torch.numel(preds)"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YHBUrkT1qUYk"
      },
      "source": [
        "def get_au_fa_fr(probs, labels):\n",
        "    sorted_probs, _ = torch.sort(probs)\n",
        "    sorted_probs = torch.cat((torch.Tensor([0]), sorted_probs, torch.Tensor([1])))\n",
        "    labels = torch.cat(labels, dim=0)\n",
        "        \n",
        "    FAs, FRs = [], []\n",
        "    for prob in sorted_probs:\n",
        "        preds = (probs >= prob) * 1\n",
        "        FA, FR = count_FA_FR(preds, labels)        \n",
        "        FAs.append(FA)\n",
        "        FRs.append(FR)\n",
        "    # plt.plot(FAs, FRs)\n",
        "    # plt.show()\n",
        "\n",
        "    # ~ area under curve using trapezoidal rule\n",
        "    return -np.trapz(FRs, x=FAs)"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CcEP5cEZqUYl"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2cP_pFIsy5p2",
        "outputId": "0b283ce6-e0ae-49d7-8386-1b4084070051"
      },
      "source": [
        "class Attention(nn.Module):\n",
        "\n",
        "    def __init__(self, hidden_size: int):\n",
        "        super().__init__()\n",
        "\n",
        "        self.energy = nn.Sequential(\n",
        "            nn.Linear(hidden_size, hidden_size),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(hidden_size, 1)\n",
        "        )\n",
        "    \n",
        "    def forward(self, input):\n",
        "        energy = self.energy(input)\n",
        "        alpha = torch.softmax(energy, dim=-2)\n",
        "        return (input * alpha).sum(dim=-2)\n",
        "\n",
        "class CRNN(nn.Module):\n",
        "\n",
        "    def __init__(self, config: TaskConfig):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                in_channels=1, out_channels=config.cnn_out_channels,\n",
        "                kernel_size=config.kernel_size, stride=config.stride\n",
        "            ),\n",
        "            nn.Flatten(start_dim=1, end_dim=2),\n",
        "        )\n",
        "\n",
        "        self.conv_out_frequency = (config.n_mels - config.kernel_size[0]) // \\\n",
        "            config.stride[0] + 1\n",
        "        \n",
        "        self.gru = nn.GRU(\n",
        "            input_size=self.conv_out_frequency * config.cnn_out_channels,\n",
        "            hidden_size=config.hidden_size,\n",
        "            num_layers=config.gru_num_layers,\n",
        "            dropout=0.1,\n",
        "            bidirectional=config.bidirectional,\n",
        "            batch_first=True\n",
        "        )\n",
        "\n",
        "        self.attention = Attention(config.hidden_size)\n",
        "        self.classifier = nn.Linear(config.hidden_size, config.num_classes)\n",
        "    \n",
        "    def forward(self, input):\n",
        "        input = input.unsqueeze(dim=1)\n",
        "        conv_output = self.conv(input).transpose(-1, -2)\n",
        "        gru_output, _ = self.gru(conv_output)\n",
        "        contex_vector = self.attention(gru_output)\n",
        "        output = self.classifier(contex_vector)\n",
        "        return output\n",
        "\n",
        "config = TaskConfig()\n",
        "model = CRNN(config)\n",
        "model"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "CRNN(\n",
              "  (conv): Sequential(\n",
              "    (0): Conv2d(1, 8, kernel_size=(5, 20), stride=(2, 8))\n",
              "    (1): Flatten(start_dim=1, end_dim=2)\n",
              "  )\n",
              "  (gru): GRU(144, 64, num_layers=2, batch_first=True, dropout=0.1)\n",
              "  (attention): Attention(\n",
              "    (energy): Sequential(\n",
              "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
              "      (1): Tanh()\n",
              "      (2): Linear(in_features=64, out_features=1, bias=True)\n",
              "    )\n",
              "  )\n",
              "  (classifier): Linear(in_features=64, out_features=2, bias=True)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DmmSFvWaqUYn"
      },
      "source": [
        "def train_epoch(model, opt, loader, log_melspec, device):\n",
        "    model.train()\n",
        "    for i, (batch, labels) in tqdm(enumerate(loader), total=len(loader)):\n",
        "        batch, labels = batch.to(device), labels.to(device)\n",
        "        batch = log_melspec(batch)\n",
        "\n",
        "        opt.zero_grad()\n",
        "\n",
        "        # run model # with autocast():\n",
        "        logits = model(batch)\n",
        "        # we need probabilities so we use softmax & CE separately\n",
        "        probs = F.softmax(logits, dim=-1)\n",
        "        loss = F.cross_entropy(logits, labels)\n",
        "\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 5)\n",
        "\n",
        "        opt.step()\n",
        "\n",
        "        # logging\n",
        "        argmax_probs = torch.argmax(probs, dim=-1)\n",
        "        FA, FR = count_FA_FR(argmax_probs, labels)\n",
        "        acc = torch.sum(argmax_probs == labels) / torch.numel(argmax_probs)\n",
        "\n",
        "    return acc"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UIeRbn4tqUYo"
      },
      "source": [
        "@torch.no_grad()\n",
        "def validation(model, loader, log_melspec, device):\n",
        "    model.eval()\n",
        "\n",
        "    val_losses, accs, FAs, FRs = [], [], [], []\n",
        "    all_probs, all_labels = [], []\n",
        "    for i, (batch, labels) in tqdm(enumerate(loader)):\n",
        "        batch, labels = batch.to(device), labels.to(device)\n",
        "        batch = log_melspec(batch)\n",
        "\n",
        "        output = model(batch)\n",
        "        # we need probabilities so we use softmax & CE separately\n",
        "        probs = F.softmax(output, dim=-1)\n",
        "        loss = F.cross_entropy(output, labels)\n",
        "\n",
        "        # logging\n",
        "        argmax_probs = torch.argmax(probs, dim=-1)\n",
        "        all_probs.append(probs[:, 1].cpu())\n",
        "        all_labels.append(labels.cpu())\n",
        "        val_losses.append(loss.item())\n",
        "        accs.append(\n",
        "            torch.sum(argmax_probs == labels).item() /  # ???\n",
        "            torch.numel(argmax_probs)\n",
        "        )\n",
        "        FA, FR = count_FA_FR(argmax_probs, labels)\n",
        "        FAs.append(FA)\n",
        "        FRs.append(FR)\n",
        "\n",
        "    # area under FA/FR curve for whole loader\n",
        "    au_fa_fr = get_au_fa_fr(torch.cat(all_probs, dim=0).cpu(), all_labels)\n",
        "    return au_fa_fr"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PpyvKwp0k3IU"
      },
      "source": [
        "from collections import defaultdict\n",
        "from IPython.display import clear_output\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "history = defaultdict(list)"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GSNW-nZCJ4Q0"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q8sVpHNoocgA",
        "outputId": "dc7e829b-bdb2-4536-9d77-71f41c0d0ad9"
      },
      "source": [
        "config = TaskConfig()\n",
        "model = CRNN(config).to(config.device)\n",
        "\n",
        "print(model)\n",
        "\n",
        "opt = torch.optim.Adam(\n",
        "    model.parameters(),\n",
        "    lr=config.learning_rate,\n",
        "    weight_decay=config.weight_decay\n",
        ")"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CRNN(\n",
            "  (conv): Sequential(\n",
            "    (0): Conv2d(1, 8, kernel_size=(5, 20), stride=(2, 8))\n",
            "    (1): Flatten(start_dim=1, end_dim=2)\n",
            "  )\n",
            "  (gru): GRU(144, 64, num_layers=2, batch_first=True, dropout=0.1)\n",
            "  (attention): Attention(\n",
            "    (energy): Sequential(\n",
            "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
            "      (1): Tanh()\n",
            "      (2): Linear(in_features=64, out_features=1, bias=True)\n",
            "    )\n",
            "  )\n",
            "  (classifier): Linear(in_features=64, out_features=2, bias=True)\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VVZgkJrMbU_x"
      },
      "source": [
        "# Download basemodel"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yKdo1K0lbYgs",
        "outputId": "7f0ad6fc-b766-483c-eac4-e011d89f8296"
      },
      "source": [
        "!gdown --id 1i0J_K45Y0Ez6Biy18r8_6gmn4qy8nekB"
      ],
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1i0J_K45Y0Ez6Biy18r8_6gmn4qy8nekB\n",
            "To: /content/basemodel_kws.pt\n",
            "\r  0% 0.00/288k [00:00<?, ?B/s]\r100% 288k/288k [00:00<00:00, 39.1MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5b0Z8USVt-dh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ff2e9d60-9173-4739-a644-78c3ee94b5d8"
      },
      "source": [
        "!gdown --id 1onkpFrJDsrhYMUfOSFD8Evg9H7Q72jAf"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Permission denied: https://drive.google.com/uc?id=1onkpFrJDsrhYMUfOSFD8Evg9H7Q72jAf\n",
            "Maybe you need to change permission over 'Anyone with the link'?\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KYWUOddeblGq"
      },
      "source": [
        "if config.device.type == 'cpu':\n",
        "    basemodel = torch.load('basemodel_kws.pt', map_location=torch.device('cpu')).to(config.device)\n",
        "else:\n",
        "    basemodel = torch.load('basemodel_kws.pt').to(config.device)"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qEIh7CRob36t",
        "outputId": "e7597001-7e27-4635-c24b-407dea8ea509"
      },
      "source": [
        "validation(basemodel, val_loader, melspec_val, config.device)"
      ],
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "102it [00:11,  8.99it/s]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1.101012528716958e-05"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-xGJF_UDrbDY"
      },
      "source": [
        "# Metrics"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q5vPatZ102x0"
      },
      "source": [
        "from thop import profile  # !pip install thop"
      ],
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SW2ak9XkrZu6"
      },
      "source": [
        "def get_size_in_megabytes(model):\n",
        "    num_params = sum([p.numel() for p in model.parameters() if p.requires_grad])\n",
        "    param_size = next(model.parameters()).element_size()\n",
        "    return (num_params * param_size) / (2 ** 20)"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I2uB_Q_6yOw2"
      },
      "source": [
        "def size_of_model(model, label=\"\"):\n",
        "    torch.save(model.state_dict(), \"temp.p\")\n",
        "    size=os.path.getsize(\"temp.p\")\n",
        "    # print(\"model: \",label,' \\t','Size (KB):', size/1e3)\n",
        "    os.remove('temp.p')\n",
        "    return size"
      ],
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YA6E3O23UrrQ"
      },
      "source": [
        "def get_metrics(model, basemodel, verbose=False):\n",
        "    model_profile = profile(model, (torch.randn(128, 40, 101).to(config.device), ))  # -> (6.0 MACs, 3.0 parameters)\n",
        "    base_profile = profile(basemodel, (torch.randn(128, 40, 101).to(config.device), ))  # -> (6.0 MACs, 3.0 parameters)\n",
        "    if verbose:\n",
        "        print(\"Base model profile :\", base_profile)\n",
        "        print(\"Model profile :\", model_profile)\n",
        "    print(\"Speed up rate :\", base_profile[0] / model_profile[0])\n",
        "    basemodel_size = size_of_model(basemodel)\n",
        "    model_size = size_of_model(model)\n",
        "    if verbose:\n",
        "        print(\"Base model size :\", basemodel_size)\n",
        "        print(\"Model size :\", model_size)\n",
        "    print(\"Compression rate :\", basemodel_size / model_size)"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BFxihRxn4HT9",
        "outputId": "b26635ce-8e36-4491-a351-7dbd8bcc19a0"
      },
      "source": [
        "get_metrics(basemodel, basemodel, True)"
      ],
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.flatten.Flatten'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.container.Sequential'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_gru() for <class 'torch.nn.modules.rnn.GRU'>.\n",
            "[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.activation.Tanh'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.Attention'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.CRNN'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.flatten.Flatten'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.container.Sequential'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_gru() for <class 'torch.nn.modules.rnn.GRU'>.\n",
            "[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.activation.Tanh'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.Attention'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.CRNN'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "Base model profile : (119527424.0, 70443.0)\n",
            "Model profile : (119527424.0, 70443.0)\n",
            "Speed up rate : 1.0\n",
            "Base model size : 288422\n",
            "Model size : 288422\n",
            "Compression rate : 1.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qMgxPC_gKdol"
      },
      "source": [
        "# Download compressed model - Dark Knowledge Distillation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hTgXr266KvNu",
        "outputId": "9d6a3b88-cc1e-4400-be4e-d08851d4c72e"
      },
      "source": [
        "!gdown --id 1UhZ0a44Gw5T2Zi2hu47KwfqF8DuhN5fg"
      ],
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1UhZ0a44Gw5T2Zi2hu47KwfqF8DuhN5fg\n",
            "To: /content/compressed_model_kws.pt\n",
            "\r  0% 0.00/34.4k [00:00<?, ?B/s]\r100% 34.4k/34.4k [00:00<00:00, 70.1MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vGLTC1CUKvNw"
      },
      "source": [
        "if config.device.type == 'cpu':\n",
        "    compressed_model = torch.load('compressed_model_kws.pt', map_location=torch.device('cpu')).to(config.device)\n",
        "else:\n",
        "    compressed_model = torch.load('compressed_model_kws.pt').to(config.device)"
      ],
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mdDWg2wZL-yF",
        "outputId": "c308cd1e-2dca-4546-915b-bb956c01c1ac"
      },
      "source": [
        "get_metrics(compressed_model, basemodel)"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.flatten.Flatten'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.container.Sequential'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_gru() for <class 'torch.nn.modules.rnn.GRU'>.\n",
            "[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.activation.Tanh'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.Attention'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.CRNN'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.flatten.Flatten'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.container.Sequential'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_gru() for <class 'torch.nn.modules.rnn.GRU'>.\n",
            "[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.activation.Tanh'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.Attention'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.CRNN'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "Speed up rate : 10.485862510386957\n",
            "Compression rate : 9.167312948954294\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8BJA3darKvNx",
        "outputId": "a5884d76-aa73-440d-c07c-91747208e24b"
      },
      "source": [
        "validation(compressed_model, val_loader, melspec_val, config.device)"
      ],
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "102it [00:10,  9.72it/s]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5.257409418968238e-05"
            ]
          },
          "metadata": {},
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "koBkoNgltc_Q"
      },
      "source": [
        "# Dark Knowledge Distillation + qint8"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xrwarlg6bRU0",
        "outputId": "9122fdfb-e42e-4203-ceed-7d7a03dc4f70"
      },
      "source": [
        "quantized_model = torch.quantization.quantize_dynamic(compressed_model, {torch.nn.Linear, torch.nn.GRU}, dtype=torch.qint8)"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/nn/quantized/dynamic/modules/rnn.py:64: UserWarning: dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.1 and num_layers=1\n",
            "  \"num_layers={}\".format(dropout, num_layers))\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hoia0UfsMw82",
        "outputId": "38fcdd7d-17e5-45fc-cf06-6a49a1012799"
      },
      "source": [
        "validation(quantized_model, val_loader, melspec_val, config.device)"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "102it [00:10,  9.41it/s]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "4.988869777817761e-05"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XoivAWwcaRh3",
        "outputId": "fa7afd4d-3dc7-4bc5-9185-00392696e91a"
      },
      "source": [
        "get_metrics(quantized_model, basemodel)"
      ],
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.flatten.Flatten'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.container.Sequential'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.quantized.dynamic.modules.rnn.PackedParameter'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.container.ModuleList'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.quantized.dynamic.modules.rnn.GRU'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.quantized.modules.linear.LinearPackedParams'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.quantized.dynamic.modules.linear.Linear'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.activation.Tanh'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.Attention'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.CRNN'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.flatten.Flatten'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.container.Sequential'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "[INFO] Register count_gru() for <class 'torch.nn.modules.rnn.GRU'>.\n",
            "[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.\n",
            "\u001b[91m[WARN] Cannot find rule for <class 'torch.nn.modules.activation.Tanh'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.Attention'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "\u001b[91m[WARN] Cannot find rule for <class '__main__.CRNN'>. Treat it as zero Macs and zero Params.\u001b[00m\n",
            "Speed up rate : 46.695069506950695\n",
            "Compression rate : 12.937783160633384\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CzRfncAteG65"
      },
      "source": [
        ""
      ],
      "execution_count": 40,
      "outputs": []
    }
  ]
}