ASR.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Installs"
      ],
      "metadata": {
        "id": "UR4qfYrVoO4v"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## wandb\n",
        "\n",
        "You will need to fetch your api key from wandb.ai"
      ],
      "metadata": {
        "id": "rd5aNaLVoR_g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install wandb -q"
      ],
      "metadata": {
        "id": "mA9qZoIDcx-h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import wandb\n",
        "wandb.login(key=\"a8b39d27e28590586b7efa4abb3acaad4e91b958\")"
      ],
      "metadata": {
        "id": "PiDduMaDIARE",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3149cbc3-1960-49f8-8db4-2f9da7f6cd23"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33makh\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# config = {\n",
        "#     'epochs': 25,\n",
        "#     'batch_size' : 1024,\n",
        "#     'learning_rate' : 0.001, # 10 epochs is recommended ONLY for the early submission - you will have to train for much longer typically.\n",
        "#     # Include other parameters as needed.\n",
        "# }"
      ],
      "metadata": {
        "id": "TmNwXomLl0vz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Levenshtein\n",
        "\n",
        "This may take a while"
      ],
      "metadata": {
        "id": "ONgAWhqdoYy-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install python-Levenshtein\n",
        "!git clone --recursive https://github.com/parlance/ctcdecode.git\n",
        "!pip install wget\n",
        "%cd ctcdecode\n",
        "!pip install .\n",
        "%cd ..\n",
        "\n",
        "!pip install torchsummaryX"
      ],
      "metadata": {
        "id": "SS7a7xeEoaV9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## imports"
      ],
      "metadata": {
        "id": "IWVONJxCobPc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import random\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torchsummaryX import summary\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence\n",
        "\n",
        "import torchaudio.transforms as tat\n",
        "\n",
        "from sklearn.metrics import accuracy_score\n",
        "import gc\n",
        "\n",
        "import zipfile\n",
        "import pandas as pd\n",
        "from tqdm import tqdm\n",
        "import os\n",
        "import datetime\n",
        "\n",
        "# imports for decoding and distance calculation\n",
        "import ctcdecode\n",
        "import Levenshtein\n",
        "from ctcdecode import CTCBeamDecoder\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "print(\"Device: \", device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "78ZTCIXoof2f",
        "outputId": "50b92120-fc5e-4e04-d2c4-b6d0e6def4ad"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Device:  cuda\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Kaggle Setup"
      ],
      "metadata": {
        "id": "gg3-yJ8tok34"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MZkG0zuiA0I1",
        "outputId": "e2219a64-c4d4-43de-8f46-4c75467204e8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Fri Nov 18 01:58:32 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   73C    P8    12W /  70W |      3MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade --force-reinstall --no-deps kaggle==1.5.8\n",
        "!mkdir /root/.kaggle\n",
        "\n",
        "with open(\"/root/.kaggle/kaggle.json\", \"w+\") as f:\n",
        "    f.write('{\"username\":\"almutwakelhassan\",\"key\":\"*****************\"}') # TODO: Put your kaggle username & key here\n",
        "\n",
        "!chmod 600 /root/.kaggle/kaggle.json"
      ],
      "metadata": {
        "id": "AdUelfGhom1m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!kaggle competitions download -c 11-785-f22-hw3p2"
      ],
      "metadata": {
        "id": "dSjBwfXeoq4B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "'''\n",
        "This will take a couple minutes, but you should see at least the following:\n",
        "11-785-f22-hw3p2.zip  ctcdecode  hw3p2\n",
        "'''\n",
        "!unzip -q 11-785-f22-hw3p2.zip\n",
        "!ls"
      ],
      "metadata": {
        "id": "_ruxWP60LCQA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Google Drive"
      ],
      "metadata": {
        "id": "R9v5ewZDMpYA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from google.colab import drive\n",
        "# drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "id": "4Cp-716IMZRd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Dataset and Dataloader"
      ],
      "metadata": {
        "id": "2ORNHnSFroP0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# ARPABET PHONEME MAPPING\n",
        "# DO NOT CHANGE\n",
        "# This overwrites the phonetics.py file.\n",
        "\n",
        "CMUdict_ARPAbet = {\n",
        "    \"\" : \" \", # BLANK TOKEN\n",
        "    \"[SIL]\": \"-\", \"NG\": \"G\", \"F\" : \"f\", \"M\" : \"m\", \"AE\": \"@\", \n",
        "    \"R\"    : \"r\", \"UW\": \"u\", \"N\" : \"n\", \"IY\": \"i\", \"AW\": \"W\", \n",
        "    \"V\"    : \"v\", \"UH\": \"U\", \"OW\": \"o\", \"AA\": \"a\", \"ER\": \"R\", \n",
        "    \"HH\"   : \"h\", \"Z\" : \"z\", \"K\" : \"k\", \"CH\": \"C\", \"W\" : \"w\", \n",
        "    \"EY\"   : \"e\", \"ZH\": \"Z\", \"T\" : \"t\", \"EH\": \"E\", \"Y\" : \"y\", \n",
        "    \"AH\"   : \"A\", \"B\" : \"b\", \"P\" : \"p\", \"TH\": \"T\", \"DH\": \"D\", \n",
        "    \"AO\"   : \"c\", \"G\" : \"g\", \"L\" : \"l\", \"JH\": \"j\", \"OY\": \"O\", \n",
        "    \"SH\"   : \"S\", \"D\" : \"d\", \"AY\": \"Y\", \"S\" : \"s\", \"IH\": \"I\",\n",
        "    \"[SOS]\": \"[SOS]\", \"[EOS]\": \"[EOS]\"}\n",
        "\n",
        "CMUdict = list(CMUdict_ARPAbet.keys())\n",
        "ARPAbet = list(CMUdict_ARPAbet.values())\n",
        "\n",
        "\n",
        "PHONEMES = CMUdict\n",
        "mapping = CMUdict_ARPAbet\n",
        "LABELS = ARPAbet"
      ],
      "metadata": {
        "id": "k0v7wHRWrqH6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# You might want to play around with the mapping as a sanity check here"
      ],
      "metadata": {
        "id": "eN2kcxwXLLBb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Train Data"
      ],
      "metadata": {
        "id": "agmNBKf4JrLV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class AudioDataset(torch.utils.data.Dataset):\n",
        "\n",
        "    # For this homework, we give you full flexibility to design your data set class.\n",
        "    # Hint: The data from HW1 is very similar to this HW\n",
        "\n",
        "    #TODO\n",
        "    def __init__(self, data_path='/content/hw3p2/train-clean-360/', partition=\"train\", limit=1): \n",
        "        '''\n",
        "        Initializes the dataset.\n",
        "\n",
        "        INPUTS: What inputs do you need here?\n",
        "        '''\n",
        "\n",
        "        # Load the directory and all files in them\n",
        "\n",
        "        self.mfcc_dir = data_path + 'mfcc/'\n",
        "        \n",
        "        self.transcript_dir = data_path + 'transcript/raw' \n",
        "\n",
        "        mfcc_files = sorted(os.listdir(self.mfcc_dir))\n",
        "        transcript_files = sorted(os.listdir(self.transcript_dir))\n",
        "\n",
        "        # mfcc_files.remove(\"processed.csv\")\n",
        "        # mfcc_files.remove(\"raw\")\n",
        "\n",
        "        self.PHONEMES = PHONEMES\n",
        "\n",
        "        num_files = len(mfcc_files)\n",
        "        if partition == \"train\":\n",
        "          mfcc_files = mfcc_files[0:int(num_files * limit)]\n",
        "          transcript_files = transcript_files[0:int(num_files * limit)]\n",
        "        elif partition == \"validate\":\n",
        "          mfcc_files = mfcc_files[int(num_files * limit) + 1:]\n",
        "          transcript_files = transcript_files[int(num_files * limit) + 1:]\n",
        "\n",
        "\n",
        "        assert len(mfcc_files) == len(transcript_files) # Making sure that we have the same no. of mfcc and transcripts\n",
        "\n",
        "        self.mfccs, self.transcripts = [], []\n",
        "          \n",
        "\n",
        "        # TODO:\n",
        "        # Iterate through mfccs and transcripts\n",
        "        for i in range(0, len(mfcc_files)):\n",
        "        #   Load a single mfcc\n",
        "            mfcc = np.load(os.path.join(self.mfcc_dir, mfcc_files[i]), allow_pickle=True)\n",
        "        #   Optionally do Cepstral Normalization of mfcc\n",
        "            mfcc = (mfcc - np.mean(mfcc, axis=0))/(np.std(mfcc, axis=0)) # edit?\n",
        "\n",
        "        #   Load the corresponding transcript\n",
        "            transcript = np.load(os.path.join(self.transcript_dir, transcript_files[i]), allow_pickle=True)[1:-1] # Remove [SOS] and [EOS] from the transcript (Is there an efficient way to do this \n",
        "            # without traversing through the transcript?)\n",
        "        #   Append each mfcc to self.mfcc, transcript to self.transcript\n",
        "            self.mfccs.append(mfcc)\n",
        "            self.transcripts.append(transcript)\n",
        "\n",
        "        #TODO\n",
        "        # WHAT SHOULD THE LENGTH OF THE DATASET BE?\n",
        "        self.length = len(self.mfccs)\n",
        "        # self.transcripts = np.vectorize(self.PHONEMES.index)(self.transcripts)\n",
        "        self.transcripts = [[self.PHONEMES.index(t) for t in line] for line in self.transcripts]\n",
        "\n",
        "        #TODO\n",
        "        # HOW CAN WE REPRESENT PHONEMES? CAN WE CREATE A MAPPING FOR THEM?\n",
        "        # HINT: TENSORS CANNOT STORE NON-NUMERICAL VALUES OR STRINGS\n",
        "\n",
        "        #TODO\n",
        "        # CREATE AN ARRAY OF ALL FEATUERS AND LABELS\n",
        "        # WHAT NORMALIZATION TECHNIQUE DID YOU USE IN HW1? CAN WE USE IT HERE?\n",
        "        '''\n",
        "        You may decide to do this in __getitem__ if you wish.\n",
        "        However, doing this here will make the __init__ function take the load of\n",
        "        loading the data, and shift it away from training.\n",
        "        '''\n",
        "       \n",
        "\n",
        "    def __len__(self):\n",
        "        \n",
        "        '''\n",
        "        TODO: What do we return here?\n",
        "        '''\n",
        "        return self.length\n",
        "\n",
        "    def __getitem__(self, ind):\n",
        "        '''\n",
        "        TODO: RETURN THE MFCC COEFFICIENTS AND ITS CORRESPONDING LABELS\n",
        "\n",
        "        If you didn't do the loading and processing of the data in __init__,\n",
        "        do that here.\n",
        "\n",
        "        Once done, return a tuple of features and labels.\n",
        "        '''\n",
        "        \n",
        "        mfcc = torch.FloatTensor(self.mfccs[ind]) # TODO\n",
        "        transcript = torch.Tensor(self.transcripts[ind]) # TODO\n",
        "        return mfcc, transcript\n",
        "\n",
        "\n",
        "    def collate_fn(self,batch):\n",
        "        '''\n",
        "        TODO:\n",
        "        1.  Extract the features and labels from 'batch'\n",
        "        2.  We will additionally need to pad both features and labels,\n",
        "            look at pytorch's docs for pad_sequence\n",
        "        3.  This is a good place to perform transforms, if you so wish. \n",
        "            Performing them on batches will speed the process up a bit.\n",
        "        4.  Return batch of features, labels, lenghts of features, \n",
        "            and lengths of labels.\n",
        "        '''\n",
        "        # batch of input mfcc coefficients\n",
        "        batch_mfcc = [i for i, j in batch]\n",
        "        # batch of output phonemes\n",
        "        batch_transcript = [j for i, j in batch]\n",
        "\n",
        "        # HINT: CHECK OUT -> pad_sequence (imported above)\n",
        "        # Also be sure to check the input format (batch_first)\n",
        "        batch_mfcc_pad = torch.nn.utils.rnn.pad_sequence(batch_mfcc, batch_first=True) # TODO\n",
        "        lengths_mfcc = [len(i) for i in batch_mfcc] # TODO \n",
        "\n",
        "        batch_transcript_pad = torch.nn.utils.rnn.pad_sequence(batch_transcript, batch_first=True) # TODO\n",
        "        lengths_transcript = [len(i) for i in batch_transcript] # TODO\n",
        "        \n",
        "\n",
        "        # You may apply some transformation, Time and Frequency masking, here in the collate function;\n",
        "        # Food for thought -> Why are we applying the transformation here and not in the __getitem__?\n",
        "        #                  -> Would we apply transformation on the validation set as well?\n",
        "        #                  -> Is the order of axes / dimensions as expected for the transform functions?\n",
        "        \n",
        "        # Return the following values: padded features, padded labels, actual length of features, actual length of the labels\n",
        "        return batch_mfcc_pad, batch_transcript_pad, torch.tensor(lengths_mfcc), torch.tensor(lengths_transcript)\n",
        "\n",
        "       "
      ],
      "metadata": {
        "id": "afd0_vlbJmr_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "v_Gn7r1vdib2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class AudioDataset2(torch.utils.data.Dataset):\n",
        "\n",
        "    # For this homework, we give you full flexibility to design your data set class.\n",
        "    # Hint: The data from HW1 is very similar to this HW\n",
        "\n",
        "    #TODO\n",
        "    def __init__(self, data_path='/content/hw3p2/train-clean-360/', partition=\"train\", limit=1): \n",
        "        '''\n",
        "        Initializes the dataset.\n",
        "\n",
        "        INPUTS: What inputs do you need here?\n",
        "        '''\n",
        "\n",
        "        # Load the directory and all files in them\n",
        "\n",
        "        self.mfcc_dir = data_path + 'mfcc/'\n",
        "        \n",
        "        self.transcript_dir = data_path + 'transcript/raw' \n",
        "\n",
        "        self.mfcc_files = sorted(os.listdir(self.mfcc_dir))\n",
        "        self.transcript_files = sorted(os.listdir(self.transcript_dir))\n",
        "\n",
        "        # mfcc_files.remove(\"processed.csv\")\n",
        "        # mfcc_files.remove(\"raw\")\n",
        "\n",
        "        self.PHONEMES = PHONEMES\n",
        "\n",
        "        # num_files = len(mfcc_files)\n",
        "        # if partition == \"train\":\n",
        "        #   mfcc_files = mfcc_files[0:int(num_files * limit)]\n",
        "        #   transcript_files = transcript_files[0:int(num_files * limit)]\n",
        "        # elif partition == \"validate\":\n",
        "        #   mfcc_files = mfcc_files[int(num_files * limit) + 1:]\n",
        "        #   transcript_files = transcript_files[int(num_files * limit) + 1:]\n",
        "\n",
        "\n",
        "        # assert len(mfcc_files) == len(transcript_files) # Making sure that we have the same no. of mfcc and transcripts\n",
        "\n",
        "        self.mfccs, self.transcripts = [], []\n",
        "          \n",
        "\n",
        "        # TODO:\n",
        "        # Iterate through mfccs and transcripts\n",
        "        # for i in range(0, len(mfcc_files)):\n",
        "        \n",
        "\n",
        "        #TODO\n",
        "        # WHAT SHOULD THE LENGTH OF THE DATASET BE?\n",
        "        self.length = len(self.mfcc_files)\n",
        "        # self.transcripts = np.vectorize(self.PHONEMES.index)(self.transcripts)\n",
        "        # self.transcripts = [[self.PHONEMES.index(t) for t in line] for line in self.transcripts]\n",
        "\n",
        "        #TODO\n",
        "        # HOW CAN WE REPRESENT PHONEMES? CAN WE CREATE A MAPPING FOR THEM?\n",
        "        # HINT: TENSORS CANNOT STORE NON-NUMERICAL VALUES OR STRINGS\n",
        "\n",
        "        #TODO\n",
        "        # CREATE AN ARRAY OF ALL FEATUERS AND LABELS\n",
        "        # WHAT NORMALIZATION TECHNIQUE DID YOU USE IN HW1? CAN WE USE IT HERE?\n",
        "        '''\n",
        "        You may decide to do this in __getitem__ if you wish.\n",
        "        However, doing this here will make the __init__ function take the load of\n",
        "        loading the data, and shift it away from training.\n",
        "        '''\n",
        "       \n",
        "\n",
        "    def __len__(self):\n",
        "        \n",
        "        '''\n",
        "        TODO: What do we return here?\n",
        "        '''\n",
        "        return self.length\n",
        "\n",
        "    def __getitem__(self, ind):\n",
        "        '''\n",
        "        TODO: RETURN THE MFCC COEFFICIENTS AND ITS CORRESPONDING LABELS\n",
        "\n",
        "        If you didn't do the loading and processing of the data in __init__,\n",
        "        do that here.\n",
        "\n",
        "        Once done, return a tuple of features and labels.\n",
        "        '''\n",
        "    #   Load a single mfcc\n",
        "        mfcc = np.load(os.path.join(self.mfcc_dir, self.mfcc_files[ind]), allow_pickle=True)\n",
        "    #   Optionally do Cepstral Normalization of mfcc\n",
        "        mfcc = (mfcc - np.mean(mfcc, axis=0))/(np.std(mfcc, axis=0)) # edit?\n",
        "\n",
        "    #   Load the corresponding transcript\n",
        "        transcript = np.load(os.path.join(self.transcript_dir, self.transcript_files[ind]), allow_pickle=True)[1:-1] # Remove [SOS] and [EOS] from the transcript (Is there an efficient way to do this \n",
        "        transcript = [self.PHONEMES.index(t) for t in transcript]\n",
        "        # without traversing through the transcript?)\n",
        "    #   Append each mfcc to self.mfcc, transcript to self.transcript\n",
        "        # self.mfccs.append(mfcc)\n",
        "        # self.transcripts.append(transcript)\n",
        "        # mfcc = torch.FloatTensor(self.mfccs[ind]) # TODO\n",
        "        mfcc = torch.FloatTensor(mfcc)\n",
        "        transcript = torch.Tensor(transcript)\n",
        "        # transcript = torch.Tensor(self.transcripts[ind]) # TODO\n",
        "        return mfcc, transcript\n",
        "\n",
        "\n",
        "    def collate_fn(self,batch):\n",
        "        '''\n",
        "        TODO:\n",
        "        1.  Extract the features and labels from 'batch'\n",
        "        2.  We will additionally need to pad both features and labels,\n",
        "            look at pytorch's docs for pad_sequence\n",
        "        3.  This is a good place to perform transforms, if you so wish. \n",
        "            Performing them on batches will speed the process up a bit.\n",
        "        4.  Return batch of features, labels, lenghts of features, \n",
        "            and lengths of labels.\n",
        "        '''\n",
        "        # batch of input mfcc coefficients\n",
        "        batch_mfcc = [i for i, j in batch]\n",
        "        # batch of output phonemes\n",
        "        batch_transcript = [j for i, j in batch]\n",
        "\n",
        "        # HINT: CHECK OUT -> pad_sequence (imported above)\n",
        "        # Also be sure to check the input format (batch_first)\n",
        "        batch_mfcc_pad = torch.nn.utils.rnn.pad_sequence(batch_mfcc, batch_first=True) # TODO\n",
        "        lengths_mfcc = [len(i) for i in batch_mfcc] # TODO \n",
        "\n",
        "        batch_transcript_pad = torch.nn.utils.rnn.pad_sequence(batch_transcript, batch_first=True) # TODO\n",
        "        lengths_transcript = [len(i) for i in batch_transcript] # TODO\n",
        "        \n",
        "\n",
        "        # You may apply some transformation, Time and Frequency masking, here in the collate function;\n",
        "        # Food for thought -> Why are we applying the transformation here and not in the __getitem__?\n",
        "        #                  -> Would we apply transformation on the validation set as well?\n",
        "        #                  -> Is the order of axes / dimensions as expected for the transform functions?\n",
        "        \n",
        "        # Return the following values: padded features, padded labels, actual length of features, actual length of the labels\n",
        "        return batch_mfcc_pad, batch_transcript_pad, torch.tensor(lengths_mfcc), torch.tensor(lengths_transcript)\n",
        "\n",
        "       "
      ],
      "metadata": {
        "id": "SqBV_rYJaH1p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Test Data"
      ],
      "metadata": {
        "id": "hqDrxeHfJw4g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Test Dataloader\n",
        "#TODO\n",
        "class AudioDatasetTest(torch.utils.data.Dataset):\n",
        "    # For this homework, we give you full flexibility to design your data set class.\n",
        "    # Hint: The data from HW1 is very similar to this HW\n",
        "\n",
        "    #TODO\n",
        "    def __init__(self, data_path='/content/hw3p2/test-clean/'): \n",
        "        '''\n",
        "        Initializes the dataset.\n",
        "\n",
        "        INPUTS: What inputs do you need here?\n",
        "        '''\n",
        "\n",
        "        # Load the directory and all files in them\n",
        "\n",
        "        self.mfcc_dir = data_path + 'mfcc/'\n",
        "        \n",
        "        # self.transcript_dir = data_path + 'transcript/raw' \n",
        "\n",
        "        mfcc_files = sorted(os.listdir(self.mfcc_dir))\n",
        "        # transcript_files = sorted(os.listdir(self.transcript_dir))\n",
        "\n",
        "        # mfcc_files.remove(\"processed.csv\")\n",
        "        # mfcc_files.remove(\"raw\")\n",
        "\n",
        "        self.PHONEMES = PHONEMES\n",
        "\n",
        "        num_files = len(mfcc_files)\n",
        "        # if partition == \"train\":\n",
        "          # mfcc_files = mfcc_files[0:int(num_files * limit)]\n",
        "          # transcript_files = transcript_files[0:int(num_files * limit)]\n",
        "        # elif partition == \"validate\":\n",
        "          # mfcc_files = mfcc_files[int(num_files * limit) + 1:]\n",
        "          # transcript_files = transcript_files[int(num_files * limit) + 1:]\n",
        "\n",
        "\n",
        "        # assert len(mfcc_files) == len(transcript_files) # Making sure that we have the same no. of mfcc and transcripts\n",
        "\n",
        "        self.mfccs, self.transcripts = [], []\n",
        "          \n",
        "\n",
        "        # TODO:\n",
        "        # Iterate through mfccs and transcripts\n",
        "        for i in range(0, len(mfcc_files)):\n",
        "        #   Load a single mfcc\n",
        "            mfcc = np.load(os.path.join(self.mfcc_dir, mfcc_files[i]), allow_pickle=True)\n",
        "        #   Optionally do Cepstral Normalization of mfcc\n",
        "            mfcc = (mfcc - np.mean(mfcc, axis=0))/(np.std(mfcc, axis=0)) # edit?\n",
        "\n",
        "        #   Load the corresponding transcript\n",
        "            # transcript = np.load(os.path.join(self.transcript_dir, transcript_files[i]), allow_pickle=True)[1:-1] # Remove [SOS] and [EOS] from the transcript (Is there an efficient way to do this \n",
        "            # without traversing through the transcript?)\n",
        "        #   Append each mfcc to self.mfcc, transcript to self.transcript\n",
        "            self.mfccs.append(mfcc)\n",
        "            # self.transcripts.append(transcript)\n",
        "\n",
        "        #TODO\n",
        "        # WHAT SHOULD THE LENGTH OF THE DATASET BE?\n",
        "        self.length = len(self.mfccs)\n",
        "        # self.transcripts = np.vectorize(self.PHONEMES.index)(self.transcripts)\n",
        "        # self.transcripts = [[self.PHONEMES.index(t) for t in line] for line in self.transcripts]\n",
        "\n",
        "        #TODO\n",
        "        # HOW CAN WE REPRESENT PHONEMES? CAN WE CREATE A MAPPING FOR THEM?\n",
        "        # HINT: TENSORS CANNOT STORE NON-NUMERICAL VALUES OR STRINGS\n",
        "\n",
        "        #TODO\n",
        "        # CREATE AN ARRAY OF ALL FEATUERS AND LABELS\n",
        "        # WHAT NORMALIZATION TECHNIQUE DID YOU USE IN HW1? CAN WE USE IT HERE?\n",
        "        '''\n",
        "        You may decide to do this in __getitem__ if you wish.\n",
        "        However, doing this here will make the __init__ function take the load of\n",
        "        loading the data, and shift it away from training.\n",
        "        '''\n",
        "       \n",
        "\n",
        "    def __len__(self):\n",
        "        \n",
        "        '''\n",
        "        TODO: What do we return here?\n",
        "        '''\n",
        "        return self.length\n",
        "\n",
        "    def __getitem__(self, ind):\n",
        "        '''\n",
        "        TODO: RETURN THE MFCC COEFFICIENTS AND ITS CORRESPONDING LABELS\n",
        "\n",
        "        If you didn't do the loading and processing of the data in __init__,\n",
        "        do that here.\n",
        "\n",
        "        Once done, return a tuple of features and labels.\n",
        "        '''\n",
        "        \n",
        "        mfcc = torch.FloatTensor(self.mfccs[ind]) # TODO\n",
        "        # transcript = torch.FloatTensor(self.transcripts[ind]) # TODO\n",
        "        return mfcc #, transcript\n",
        "\n",
        "\n",
        "    def collate_fn(self,batch):\n",
        "        '''\n",
        "        TODO:\n",
        "        1.  Extract the features and labels from 'batch'\n",
        "        2.  We will additionally need to pad both features and labels,\n",
        "            look at pytorch's docs for pad_sequence\n",
        "        3.  This is a good place to perform transforms, if you so wish. \n",
        "            Performing them on batches will speed the process up a bit.\n",
        "        4.  Return batch of features, labels, lenghts of features, \n",
        "            and lengths of labels.\n",
        "        '''\n",
        "        # batch of input mfcc coefficients\n",
        "        batch_mfcc = batch\n",
        "        # batch of output phonemes\n",
        "        # batch_transcript = [j for i, j in batch]\n",
        "\n",
        "        # HINT: CHECK OUT -> pad_sequence (imported above)\n",
        "        # Also be sure to check the input format (batch_first)\n",
        "        batch_mfcc_pad = torch.nn.utils.rnn.pad_sequence(batch_mfcc, batch_first=True) # TODO\n",
        "        lengths_mfcc = [len(i) for i in batch_mfcc] # TODO \n",
        "\n",
        "        # batch_transcript_pad = torch.nn.utils.rnn.pad_sequence(batch_transcript, batch_first=True) # TODO\n",
        "        # lengths_transcript = [len(i) for i in batch_transcript_pad] # TODO\n",
        "        \n",
        "\n",
        "        # You may apply some transformation, Time and Frequency masking, here in the collate function;\n",
        "        # Food for thought -> Why are we applying the transformation here and not in the __getitem__?\n",
        "        #                  -> Would we apply transformation on the validation set as well?\n",
        "        #                  -> Is the order of axes / dimensions as expected for the transform functions?\n",
        "        \n",
        "        # Return the following values: padded features, padded labels, actual length of features, actual length of the labels\n",
        "        return batch_mfcc_pad, torch.tensor(lengths_mfcc)\n",
        "\n",
        "       "
      ],
      "metadata": {
        "id": "HrLS1wfVJppA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data - Hyperparameters"
      ],
      "metadata": {
        "id": "Pt-veYcdL6Fe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "BATCH_SIZE = 64 # Increase if your device can handle it\n",
        "\n",
        "transforms = [] # set of tranformations\n",
        "# You may pass this as a parameter to the dataset class above\n",
        "# This will help modularize your implementation\n",
        "\n",
        "root = '/content/hw3p2' "
      ],
      "metadata": {
        "id": "4icymeX1ImUN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data loaders"
      ],
      "metadata": {
        "id": "NmuPk9J6L8dz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# get me RAMMM!!!! \n",
        "import gc \n",
        "gc.collect()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3_kG0gU2x4hH",
        "outputId": "364dc5bf-22f2-4859-fd04-596f01c76270"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "253"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Create objects for the dataset class\n",
        "train_data = AudioDataset() #TODO\n",
        "val_data = AudioDataset(data_path='/content/hw3p2/dev-clean/') # TODO : You can either use the same class with some modifications or make a new one :)\n",
        "test_data = AudioDatasetTest() #TODO\n",
        "\n",
        "# Do NOT forget to pass in the collate function as parameter while creating the dataloader\n",
        "train_loader = torch.utils.data.DataLoader(train_data, num_workers= 4,\n",
        "                                           batch_size=BATCH_SIZE, pin_memory= True,\n",
        "                                           shuffle= True, collate_fn= train_data.collate_fn)\n",
        "\n",
        "val_loader = torch.utils.data.DataLoader(val_data, num_workers= 2,\n",
        "                                         batch_size=BATCH_SIZE, pin_memory= True,\n",
        "                                         shuffle= False, collate_fn= val_data.collate_fn)\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(test_data, num_workers= 2, \n",
        "                                          batch_size=BATCH_SIZE, pin_memory= True, \n",
        "                                          shuffle= False, collate_fn= test_data.collate_fn)\n",
        "\n",
        "print(\"Batch size: \", BATCH_SIZE)\n",
        "print(\"Train dataset samples = {}, batches = {}\".format(train_data.__len__(), len(train_loader)))\n",
        "print(\"Val dataset samples = {}, batches = {}\".format(val_data.__len__(), len(val_loader)))\n",
        "print(\"Test dataset samples = {}, batches = {}\".format(test_data.__len__(), len(test_loader)))"
      ],
      "metadata": {
        "id": "4mzoYfTKu14s",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cd8c7d20-afe0-4fa2-d789-d8399211d457"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch size:  64\n",
            "Train dataset samples = 104014, batches = 1626\n",
            "Val dataset samples = 2703, batches = 43\n",
            "Test dataset samples = 2620, batches = 41\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# sanity check\n",
        "for data in train_loader:\n",
        "    x, y, lx, ly = data\n",
        "    print(x.shape, y.shape, lx.shape, ly.shape)\n",
        "    break \n",
        "# torch.Size([64, 1678, 15]) torch.Size([64, 212]) torch.Size([64]) torch.Size([64])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xg7UzWM3SNUO",
        "outputId": "02c4dc58-c163-451b-cdf7-b47277360d63"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([64, 1679, 15]) torch.Size([64, 197]) torch.Size([64]) torch.Size([64])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# sanity check\n",
        "for data in train_loader:\n",
        "    x, y, lx, ly = data\n",
        "    print(x.shape, y.shape, lx.shape, ly.shape)\n",
        "    break "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cXMtwyviKaxK",
        "outputId": "7109b4fe-a524-4535-b52f-36fcc9a80b8d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([64, 1691, 15]) torch.Size([64, 198]) torch.Size([64]) torch.Size([64])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model Config"
      ],
      "metadata": {
        "id": "Ly4mjUUUuJhy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "OUT_SIZE = len(LABELS)\n",
        "OUT_SIZE"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RZ-qQ_Sf-LIu",
        "outputId": "05cf819e-e5f0-4d02-90a9-962672da1f07"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "43"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Basic"
      ],
      "metadata": {
        "id": "HLad4pChcuvX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "\n",
        "class Network(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Network, self).__init__()\n",
        "        # Adding some sort of embedding layer or feature extractor might help performance.\n",
        "        # self.embedding = ?\n",
        "        # TODO : look up the documentation. You might need to pass some additional parameters.\n",
        "        self.lstm = nn.LSTM(input_size = 15, hidden_size = 256, num_layers = 1, batch_first=True, bidirectional=True) \n",
        "        # hidden_size = 256\n",
        "\n",
        "        self.classification = nn.Sequential(\n",
        "            #TODO: Linear layer with in_features from the lstm module above and out_features = OUT_SIZE\n",
        "            torch.nn.Linear(512, OUT_SIZE)\n",
        "        )        \n",
        "        self.logSoftmax = torch.nn.LogSoftmax(dim=2)\n",
        "        #TODO: Apply a log softmax here. Which dimension would apply it on ?\n",
        "\n",
        "    def forward(self, x, lx):\n",
        "        #TODO\n",
        "        # The forward function takes 2 parameter inputs here. Why?\n",
        "        # Refer to the handout for hints\n",
        "        packed = torch.nn.utils.rnn.pack_padded_sequence(x, lx, batch_first=True, enforce_sorted=False)\n",
        "        packed, (x, y) = self.lstm.forward(packed)\n",
        "        seq_unpacked, lens_unpacked = torch.nn.utils.rnn.pad_packed_sequence(packed, batch_first=True)\n",
        "        z = self.classification.forward(seq_unpacked)\n",
        "        z = self.logSoftmax.forward(z)\n",
        "        return z, lens_unpacked\n"
      ],
      "metadata": {
        "id": "EQhvHr71GJfq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "\n",
        "class CNN_LSTM_Model(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(CNN_LSTM_Model, self).__init__()\n",
        "        # Adding some sort of embedding layer or feature extractor might help performance.\n",
        "        # self.embedding = ?\n",
        "        # TODO : look up the documentation. You might need to pass some additional parameters.\n",
        "        self.feature_extraction = torch.nn.Sequential(\n",
        "            torch.nn.Conv1d(15, 256, kernel_size=3, padding=1),\n",
        "            torch.nn.BatchNorm1d(256),\n",
        "            torch.nn.GELU(),\n",
        "\n",
        "            torch.nn.Conv1d(256, 512, kernel_size=3, padding=1),\n",
        "            torch.nn.BatchNorm1d(512),\n",
        "            torch.nn.GELU(),\n",
        "\n",
        "            torch.nn.Conv1d(512, 512, kernel_size=3, padding=1),\n",
        "            torch.nn.AvgPool1d(kernel_size=3, stride=2)\n",
        "        )\n",
        "        self.lstm = nn.LSTM(input_size = 512, hidden_size = 1024, num_layers = 2, \n",
        "                            dropout=0.5, batch_first=True, bidirectional=True) \n",
        "        # self.lstm2 = nn.LSTM(input_size = 512, hidden_size = 256, num_layers = 1, batch_first=True, bidirectional=True) \n",
        "\n",
        "\n",
        "        self.classification = nn.Sequential(\n",
        "            #TODO: Linear layer with in_features from the lstm module above and out_features = OUT_SIZE\n",
        "            \n",
        "            torch.nn.Linear(2048, 2048),\n",
        "            torch.nn.GELU(),\n",
        "            torch.nn.Dropout(0.3),\n",
        "            \n",
        "            torch.nn.Linear(2048, 1024),\n",
        "            torch.nn.GELU(),\n",
        "            torch.nn.Dropout(0.3),\n",
        "\n",
        "            torch.nn.Linear(1024, 512),\n",
        "            torch.nn.GELU(),\n",
        "            torch.nn.Linear(512, OUT_SIZE)\n",
        "\n",
        "        )        \n",
        "        self.logSoftmax = torch.nn.LogSoftmax(dim=2)\n",
        "        #TODO: Apply a log softmax here. Which dimension would apply it on ?\n",
        "\n",
        "    def forward(self, x, lx):\n",
        "        #TODO\n",
        "        # The forward function takes 2 parameter inputs here. Why?\n",
        "        # Refer to the handout for hints\n",
        "        extraction = self.feature_extraction(x.permute(0, 2, 1))\n",
        "        extraction = extraction.permute(0, 2, 1)\n",
        "\n",
        "        lx = lx.clamp(max=extraction.shape[1])\n",
        "\n",
        "        packed = torch.nn.utils.rnn.pack_padded_sequence(extraction, lx, batch_first=True, enforce_sorted=False)\n",
        "        packed1, (x, y) = self.lstm(packed)\n",
        "        seq_unpacked, lens_unpacked = torch.nn.utils.rnn.pad_packed_sequence(packed1, batch_first=True)\n",
        "        z = self.classification(seq_unpacked)\n",
        "        z = self.logSoftmax(z)\n",
        "        \n",
        "        return z, lens_unpacked\n"
      ],
      "metadata": {
        "id": "aDKvatNfznRs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## INIT"
      ],
      "metadata": {
        "id": "tUThsowyQdN7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "model = CNN_LSTM_Model().to(device)\n",
        "summary(model, x.to(device), lx) # x and lx come from the sanity check above :)"
      ],
      "metadata": {
        "id": "CGoiXd70tb5z",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "30a3ef8f-81bd-4213-f8fe-1c2aa04d85a9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============================================================================================\n",
            "                                     Kernel Shape     Output Shape  \\\n",
            "Layer                                                                \n",
            "0_feature_extraction.Conv1d_0        [15, 256, 3]  [64, 256, 1691]   \n",
            "1_feature_extraction.BatchNorm1d_1          [256]  [64, 256, 1691]   \n",
            "2_feature_extraction.GELU_2                     -  [64, 256, 1691]   \n",
            "3_feature_extraction.Conv1d_3       [256, 512, 3]  [64, 512, 1691]   \n",
            "4_feature_extraction.BatchNorm1d_4          [512]  [64, 512, 1691]   \n",
            "5_feature_extraction.GELU_5                     -  [64, 512, 1691]   \n",
            "6_feature_extraction.Conv1d_6       [512, 512, 3]  [64, 512, 1691]   \n",
            "7_feature_extraction.AvgPool1d_7                -   [64, 512, 845]   \n",
            "8_lstm                                          -    [52953, 2048]   \n",
            "9_classification.Linear_0            [2048, 2048]  [64, 845, 2048]   \n",
            "10_classification.GELU_1                        -  [64, 845, 2048]   \n",
            "11_classification.Dropout_2                     -  [64, 845, 2048]   \n",
            "12_classification.Linear_3           [2048, 1024]  [64, 845, 1024]   \n",
            "13_classification.GELU_4                        -  [64, 845, 1024]   \n",
            "14_classification.Dropout_5                     -  [64, 845, 1024]   \n",
            "15_classification.Linear_6            [1024, 512]   [64, 845, 512]   \n",
            "16_classification.GELU_7                        -   [64, 845, 512]   \n",
            "17_classification.Linear_8              [512, 43]    [64, 845, 43]   \n",
            "18_logSoftmax                                   -    [64, 845, 43]   \n",
            "\n",
            "                                        Params     Mult-Adds  \n",
            "Layer                                                         \n",
            "0_feature_extraction.Conv1d_0          11.776k     19.48032M  \n",
            "1_feature_extraction.BatchNorm1d_1       512.0         256.0  \n",
            "2_feature_extraction.GELU_2                  -             -  \n",
            "3_feature_extraction.Conv1d_3         393.728k   664.928256M  \n",
            "4_feature_extraction.BatchNorm1d_4      1.024k         512.0  \n",
            "5_feature_extraction.GELU_5                  -             -  \n",
            "6_feature_extraction.Conv1d_6         786.944k  1.329856512G  \n",
            "7_feature_extraction.AvgPool1d_7             -             -  \n",
            "8_lstm                              37.781504M    37.748736M  \n",
            "9_classification.Linear_0            4.196352M     4.194304M  \n",
            "10_classification.GELU_1                     -             -  \n",
            "11_classification.Dropout_2                  -             -  \n",
            "12_classification.Linear_3           2.098176M     2.097152M  \n",
            "13_classification.GELU_4                     -             -  \n",
            "14_classification.Dropout_5                  -             -  \n",
            "15_classification.Linear_6              524.8k      524.288k  \n",
            "16_classification.GELU_7                     -             -  \n",
            "17_classification.Linear_8             22.059k       22.016k  \n",
            "18_logSoftmax                                -             -  \n",
            "--------------------------------------------------------------------------------------------\n",
            "                            Totals\n",
            "Total params            45.816875M\n",
            "Trainable params        45.816875M\n",
            "Non-trainable params           0.0\n",
            "Mult-Adds             2.058852352G\n",
            "============================================================================================\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\n",
              "  <div id=\"df-d5a813ea-8ffd-455b-82cc-ede3d11adcd3\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <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>Kernel Shape</th>\n",
              "      <th>Output Shape</th>\n",
              "      <th>Params</th>\n",
              "      <th>Mult-Adds</th>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>Layer</th>\n",
              "      <th></th>\n",
              "      <th></th>\n",
              "      <th></th>\n",
              "      <th></th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0_feature_extraction.Conv1d_0</th>\n",
              "      <td>[15, 256, 3]</td>\n",
              "      <td>[64, 256, 1691]</td>\n",
              "      <td>11776.0</td>\n",
              "      <td>1.948032e+07</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1_feature_extraction.BatchNorm1d_1</th>\n",
              "      <td>[256]</td>\n",
              "      <td>[64, 256, 1691]</td>\n",
              "      <td>512.0</td>\n",
              "      <td>2.560000e+02</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2_feature_extraction.GELU_2</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 256, 1691]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3_feature_extraction.Conv1d_3</th>\n",
              "      <td>[256, 512, 3]</td>\n",
              "      <td>[64, 512, 1691]</td>\n",
              "      <td>393728.0</td>\n",
              "      <td>6.649283e+08</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4_feature_extraction.BatchNorm1d_4</th>\n",
              "      <td>[512]</td>\n",
              "      <td>[64, 512, 1691]</td>\n",
              "      <td>1024.0</td>\n",
              "      <td>5.120000e+02</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5_feature_extraction.GELU_5</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 512, 1691]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>6_feature_extraction.Conv1d_6</th>\n",
              "      <td>[512, 512, 3]</td>\n",
              "      <td>[64, 512, 1691]</td>\n",
              "      <td>786944.0</td>\n",
              "      <td>1.329857e+09</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7_feature_extraction.AvgPool1d_7</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 512, 845]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>8_lstm</th>\n",
              "      <td>-</td>\n",
              "      <td>[52953, 2048]</td>\n",
              "      <td>37781504.0</td>\n",
              "      <td>3.774874e+07</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>9_classification.Linear_0</th>\n",
              "      <td>[2048, 2048]</td>\n",
              "      <td>[64, 845, 2048]</td>\n",
              "      <td>4196352.0</td>\n",
              "      <td>4.194304e+06</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>10_classification.GELU_1</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 845, 2048]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>11_classification.Dropout_2</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 845, 2048]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>12_classification.Linear_3</th>\n",
              "      <td>[2048, 1024]</td>\n",
              "      <td>[64, 845, 1024]</td>\n",
              "      <td>2098176.0</td>\n",
              "      <td>2.097152e+06</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>13_classification.GELU_4</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 845, 1024]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>14_classification.Dropout_5</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 845, 1024]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>15_classification.Linear_6</th>\n",
              "      <td>[1024, 512]</td>\n",
              "      <td>[64, 845, 512]</td>\n",
              "      <td>524800.0</td>\n",
              "      <td>5.242880e+05</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>16_classification.GELU_7</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 845, 512]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>17_classification.Linear_8</th>\n",
              "      <td>[512, 43]</td>\n",
              "      <td>[64, 845, 43]</td>\n",
              "      <td>22059.0</td>\n",
              "      <td>2.201600e+04</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>18_logSoftmax</th>\n",
              "      <td>-</td>\n",
              "      <td>[64, 845, 43]</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-d5a813ea-8ffd-455b-82cc-ede3d11adcd3')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-d5a813ea-8ffd-455b-82cc-ede3d11adcd3 button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-d5a813ea-8ffd-455b-82cc-ede3d11adcd3');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ],
            "text/plain": [
              "                                     Kernel Shape  ...     Mult-Adds\n",
              "Layer                                              ...              \n",
              "0_feature_extraction.Conv1d_0        [15, 256, 3]  ...  1.948032e+07\n",
              "1_feature_extraction.BatchNorm1d_1          [256]  ...  2.560000e+02\n",
              "2_feature_extraction.GELU_2                     -  ...           NaN\n",
              "3_feature_extraction.Conv1d_3       [256, 512, 3]  ...  6.649283e+08\n",
              "4_feature_extraction.BatchNorm1d_4          [512]  ...  5.120000e+02\n",
              "5_feature_extraction.GELU_5                     -  ...           NaN\n",
              "6_feature_extraction.Conv1d_6       [512, 512, 3]  ...  1.329857e+09\n",
              "7_feature_extraction.AvgPool1d_7                -  ...           NaN\n",
              "8_lstm                                          -  ...  3.774874e+07\n",
              "9_classification.Linear_0            [2048, 2048]  ...  4.194304e+06\n",
              "10_classification.GELU_1                        -  ...           NaN\n",
              "11_classification.Dropout_2                     -  ...           NaN\n",
              "12_classification.Linear_3           [2048, 1024]  ...  2.097152e+06\n",
              "13_classification.GELU_4                        -  ...           NaN\n",
              "14_classification.Dropout_5                     -  ...           NaN\n",
              "15_classification.Linear_6            [1024, 512]  ...  5.242880e+05\n",
              "16_classification.GELU_7                        -  ...           NaN\n",
              "17_classification.Linear_8              [512, 43]  ...  2.201600e+04\n",
              "18_logSoftmax                                   -  ...           NaN\n",
              "\n",
              "[19 rows x 4 columns]"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training Config"
      ],
      "metadata": {
        "id": "IBwunYpyugFg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_config = {\n",
        "    \"beam_width\" : 3,\n",
        "    \"lr\" : 1e-3,\n",
        "    \"epochs\" : 100\n",
        "    } # Feel free to add more items here"
      ],
      "metadata": {
        "id": "SBRz7gCQqzGE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "run = wandb.init(\n",
        "    name = \"CNN_LSTM_Separate002\", ## Wandb creates random run names if you skip this field\n",
        "    reinit = True, ### Allows reinitalizing runs when you re-run this cell\n",
        "    # run_id = ### Insert specific run id here if you want to resume a previous run\n",
        "    # resume = \"must\" ### You need this to resume previous runs, but comment out reinit = True when using this\n",
        "    project = \"hw3p2-ablations\", ### Project should be created in your wandb account \n",
        "    config = train_config ### Wandb Config for your run\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "id": "HxOod7PIFIMk",
        "outputId": "8f47e1b2-39e3-48ee-f1d4-b0866913891f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "Tracking run with wandb version 0.13.5"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "Run data is saved locally in <code>/content/wandb/run-20221118_020101-2qec087r</code>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "Syncing run <strong><a href=\"https://wandb.ai/akh/hw3p2-ablations/runs/2qec087r\" target=\"_blank\">CNN_LSTM_Separate002</a></strong> to <a href=\"https://wandb.ai/akh/hw3p2-ablations\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_config"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OGB9c4nut7ir",
        "outputId": "ce974ce4-270d-4683-f941-c8f49214c454"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'beam_width': 3, 'epochs': 100, 'lr': 0.001}"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#TODO\n",
        "criterion = torch.nn.CTCLoss(reduction = 'mean')\n",
        "# Define CTC loss as the criterion. How would the losses be reduced?\n",
        "# CTC Loss: https://pytorch.org/docs/stable/generated/torch.nn.CTCLoss.html\n",
        "# Refer to the handout for hints\n",
        "\n",
        "optimizer =  torch.optim.AdamW(model.parameters(), lr=train_config['lr']) # What goes in here?\n",
        "\n",
        "# Declare the decoder. Use the CTC Beam Decoder to decode phonemes\n",
        "# CTC Beam Decoder Doc: https://github.com/parlance/ctcdecode\n",
        "decoder = ctcdecode.CTCBeamDecoder(\n",
        "    labels = ARPAbet,\n",
        "    model_path=None,\n",
        "    alpha=0,\n",
        "    beta=0,\n",
        "    cutoff_top_n=40,\n",
        "    cutoff_prob=1.0,\n",
        "    beam_width=train_config['beam_width'],\n",
        "    num_processes=4,\n",
        "    blank_id=0,\n",
        "    log_probs_input=True)\n",
        "\n",
        "# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, verbose=True, factor=0.25)\n",
        "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, train_config[\"epochs\"] * len(train_loader))\n",
        "\n",
        "# Mixed Precision, if you need it\n",
        "scaler = torch.cuda.amp.GradScaler()"
      ],
      "metadata": {
        "id": "iGoozH2nd6KB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Levenshtein"
      ],
      "metadata": {
        "id": "Jmc6_4eWL2Xp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Use debug = True to see debug outputs\n",
        "def calculate_levenshtein(h, y, lh, ly, decoder, labels, debug = False):\n",
        "\n",
        "    if debug:\n",
        "        pass\n",
        "        # print(f\"\\n----- IN LEVENSHTEIN -----\\n\")\n",
        "        # Add any other debug statements as you may need\n",
        "        # you may want to use debug in several places in this function\n",
        "        \n",
        "    # TODO: look at docs for CTC.decoder and find out what is returned here\n",
        "    beam_results, beam_scores, timesteps, out_lens = decoder.decode(h, seq_lens = lh)\n",
        "    # print(beam_results, beam_scores, timesteps, out_lens)\n",
        "    batch_size = h.shape[0] # TODO\n",
        "    distance = 0 # Initialize the distance to be 0 initially\n",
        "\n",
        "    for i in range(batch_size): \n",
        "        # TODO: Loop through each element in the batch\n",
        "        h_slice = beam_results[i, 0, :out_lens[i, 0]]\n",
        "        h_string = [labels[x] for x in h_slice]\n",
        "        y_slice = y[i, :ly[i]]\n",
        "        \n",
        "        y_string = [labels[int(x)] for x in y_slice]\n",
        "\n",
        "        distance += Levenshtein.distance(\"\".join(h_string), \"\".join(y_string))\n",
        "\n",
        "        del h_slice, h_string, y_slice, y_string\n",
        "    del beam_results, beam_scores, timesteps, out_lens\n",
        "\n",
        "    distance /= batch_size # TODO: Uncomment this, but think about why we are doing this\n",
        "\n",
        "    return distance"
      ],
      "metadata": {
        "id": "KHjnCDddL36E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# ANOTEHR SANITY CHECK\n",
        "\n",
        "with torch.no_grad():\n",
        "  for i, data in enumerate(train_loader):\n",
        "      \n",
        "      #TODO: \n",
        "      # Follow the following steps, and \n",
        "      # Add some print statements here for sanity checking\n",
        "      # print(data)\n",
        "      h, y, lh, ly = data\n",
        "      h = h.to(device)\n",
        "      y = y.to(device)\n",
        "      # lh.to(device)\n",
        "      # ly.to(device)\n",
        "      # print(\"h\", h.shape)\n",
        "      # print(\"y\", y.shape)\n",
        "      # print(\"lh\", lh.shape)\n",
        "      # print(\"ly\", ly.shape)\n",
        "\n",
        "      out, out_len = model(h, lh)\n",
        "      # print(\"out\", out.shape)\n",
        "      # print(\"out_len\", out_len.shape)\n",
        "      distance = calculate_levenshtein(out, y, out_len, ly, decoder, LABELS)\n",
        "      loss = criterion(out.permute(1, 0, 2), y, out_len, ly)\n",
        "\n",
        "      # print(make_output(out, out_len, decoder, LABELS))\n",
        "      # print(make_output(y, ly, decoder, LABELS))\n",
        "      #1. What values are you returning from the collate function\n",
        "      #2. Move the features and target to <DEVICE>\n",
        "      #3. Print the shapes of each to get a fair understanding \n",
        "      #4. Pass the inputs to the model\n",
        "            # Think of the following before you implement:\n",
        "            # 4.1 What will be the input to your model?\n",
        "            # 4.2 What would the model output?\n",
        "            # 4.3 Print the shapes of the output to get a fair understanding \n",
        "\n",
        "      # Calculate loss: https://pytorch.org/docs/stable/generated/torch.nn.CTCLoss.html\n",
        "      # Calculating the loss is not straightforward. Check the input format of each parameter\n",
        "\n",
        "      print(f\"lev-distance: {distance}\")\n",
        "      print(f\"loss: {loss}\")\n",
        "\n",
        "\n",
        "      break # one iteration is enough"
      ],
      "metadata": {
        "id": "GnTLL-5gMBrY",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "34a63953-1394-4451-df30-34cf16784f75"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "lev-distance: 301.546875\n",
            "loss: 20.174972534179688\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "6fLLj5KIMMOe"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Eval function\n",
        "Writing a function to do one round of evaluations will help make your code more modular, you can, however, choose to skip this if you'd like it."
      ],
      "metadata": {
        "id": "kH0RAbCaMl9a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(data_loader, model):\n",
        "    model.eval()\n",
        "    dist = 0\n",
        "    loss = 0\n",
        "    batch_bar = tqdm(total=len(data_loader), dynamic_ncols=True, leave=False, position=0, desc='Val') \n",
        "    # TODO Fill this function out, if you're using it.\n",
        "    ## Forgot this\n",
        "    with torch.no_grad():\n",
        "      for i, data in enumerate(data_loader):\n",
        "\n",
        "          # TODO: Fill this with the help of your sanity check\n",
        "          h, y, lh, ly = data\n",
        "          h = h.to(device)\n",
        "          y = y.to(device)\n",
        "          # ly = ly.to(device)\n",
        "\n",
        "          out, out_len = model(h, lh)\n",
        "          loss += criterion(out.permute(1, 0, 2), y, out_len, ly)\n",
        "          dist += calculate_levenshtein(out, y, out_len, ly, decoder, LABELS)\n",
        "\n",
        "          del h, y, lh, ly, data, out, out_len\n",
        "          torch.cuda.empty_cache()\n",
        "          # torch.clear_autocast_cache()\n",
        "          gc.collect()\n",
        "\n",
        "          batch_bar.set_postfix(\n",
        "              loss = f\"{loss/ (i+1):.4f}\",\n",
        "              dist = f\"{dist/(i+1):.4f}\",\n",
        "              lr = f\"{optimizer.param_groups[0]['lr']}\"\n",
        "          )\n",
        "          batch_bar.update()\n",
        "      batch_bar.close()\n",
        "\n",
        "    return dist/len(data_loader) ,loss/len(data_loader),"
      ],
      "metadata": {
        "id": "0nqLiAmkMMBc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training Setup"
      ],
      "metadata": {
        "id": "qpYExu4vT4_g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# This is for checkpointing, if you're doing it over multiple sessions\n",
        "\n",
        "last_epoch_completed = 0\n",
        "start = last_epoch_completed\n",
        "end = train_config['epochs'] #epochs\n",
        "best_val_dist = float(\"inf\") # if you're restarting from some checkpoint, use what you saw there.\n",
        "dist_freq = 1"
      ],
      "metadata": {
        "id": "tExvyl1BIdMC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Again, writing a train step might help you code be more modular. You may choose to skip this and write the whole thing out in the training loop below if you so wish."
      ],
      "metadata": {
        "id": "pGn17rLw9ChF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train_step(train_loader, model, optimizer, criterion, scheduler, scaler):\n",
        "    \n",
        "    model.train()\n",
        "    batch_bar = tqdm(total=len(train_loader), dynamic_ncols=True, leave=False, position=0, desc='Train') \n",
        "    train_loss = 0\n",
        "\n",
        "    for i, data in enumerate(train_loader):\n",
        "\n",
        "        # TODO: Fill this with the help of your sanity check\n",
        "\n",
        "        h, y, lh, ly = data\n",
        "        h = h.to(device)\n",
        "        y = y.to(device)\n",
        "        # lh.to(device)\n",
        "        # ly.to(device)\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # HINT: Are you using mixed precision?\n",
        "        \n",
        "        with torch.cuda.amp.autocast():\n",
        "            out, out_len = model.forward(h, lh)\n",
        "            loss = criterion(out.permute(1, 0, 2), y, out_len, ly)\n",
        "        # loss.backward()\n",
        "        \n",
        "        # print(loss.item())\n",
        "        # h.to('cpu')\n",
        "        # y.to('cpu')\n",
        "        # ly.to('cpu')\n",
        "        \n",
        "\n",
        "        batch_bar.set_postfix(\n",
        "            tloss = f\"{train_loss/ (i+1):.4f}\",\n",
        "            lr = f\"{optimizer.param_groups[0]['lr']}\"\n",
        "        )\n",
        "\n",
        "        train_loss += loss\n",
        "        scaler.scale(loss).backward()\n",
        "        scaler.step(optimizer)\n",
        "        scaler.update()\n",
        "        scheduler.step()\n",
        "        # optimizer.step()\n",
        "        batch_bar.update()\n",
        "        \n",
        "        del h, y, lh, ly\n",
        "    \n",
        "    batch_bar.close()\n",
        "    train_loss /= len(train_loader) # TODO\n",
        "\n",
        "    return train_loss # And anything else you may wish to get out of this function"
      ],
      "metadata": {
        "id": "_vH4QStLUjH8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Train Loop"
      ],
      "metadata": {
        "id": "MY69hgxUXhTI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "#TODO: Please complete the training loop\n",
        "\n",
        "for epoch in range(train_config[\"epochs\"]):\n",
        "\n",
        "    # one training step\n",
        "    # one validation step (if you want)\n",
        "\n",
        "    # HINT: Calculating levenshtein distance takes a long time. Do you need to do it every epoch?\n",
        "    # Does the training step even need it? \n",
        "\n",
        "    # Where you have your scheduler.step depends on the scheduler you use.\n",
        "    train_loss = train_step(train_loader, model, optimizer, criterion, scheduler, scaler)\n",
        "    val_dist, val_loss = evaluate(val_loader, model)\n",
        "    print(f\"\\nEpoch {epoch} / Val Distance {val_dist} / Train Loss {train_loss}\")\n",
        "    # Use the below code to save models\n",
        "    if val_dist < best_val_dist:\n",
        "      #path = os.path.join(root_path, model_directory, 'checkpoint' + '.pth')\n",
        "      print(\"Saving model\")\n",
        "      torch.save({'model_state_dict':model.state_dict(),\n",
        "                  'optimizer_state_dict':optimizer.state_dict(),\n",
        "                  'val_dist': val_dist, \n",
        "                  'epoch': epoch}, './checkpoint.pth')\n",
        "      best_val_dist = val_dist\n",
        "      # wandb.save('checkpoint.pth')\n",
        "\n",
        "    # scheduler.step(val_dist)\n",
        "    \n",
        "    curr_lr = float(optimizer.param_groups[0]['lr'])\n",
        "    # You may want to log some hyperparameters and results on wandb\n",
        "    # wandb.log({\"train_loss\":train_loss, 'val_dist': val_dist, \"learning_Rate\": curr_lr})\n",
        "\n",
        "run.finish()"
      ],
      "metadata": {
        "id": "JR43E28rM9Ak",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "31fc416f-7920-4eb7-be73-627fe75579fd"
      },
      "execution_count": null,
      "outputs": [
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 0 / Val Distance 11.109883720930233 / Train Loss 1.3558909893035889\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 1 / Val Distance 8.447625968992247 / Train Loss 0.43633466958999634\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 2 / Val Distance 7.211627906976744 / Train Loss 0.35278111696243286\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 3 / Val Distance 6.621608527131784 / Train Loss 0.3165063261985779\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 4 / Val Distance 6.0418604651162795 / Train Loss 0.2856194078922272\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 5 / Val Distance 5.954312015503876 / Train Loss 0.26404061913490295\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 6 / Val Distance 5.547407945736435 / Train Loss 0.25072410702705383\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 7 / Val Distance 5.541497093023256 / Train Loss 0.23604685068130493\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 8 / Val Distance 5.378415697674418 / Train Loss 0.22991125285625458\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 9 / Val Distance 5.245930232558139 / Train Loss 0.2205376774072647\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 10 / Val Distance 5.1955184108527135 / Train Loss 0.21563449501991272\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 11 / Val Distance 4.971027131782946 / Train Loss 0.20572160184383392\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 12 / Val Distance 5.238081395348837 / Train Loss 0.2013939768075943\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 13 / Val Distance 4.797795542635659 / Train Loss 0.19594600796699524\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 14 / Val Distance 4.782000968992248 / Train Loss 0.18980638682842255\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 15 / Val Distance 4.765285852713179 / Train Loss 0.1848566234111786\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 16 / Val Distance 4.718459302325582 / Train Loss 0.17866967618465424\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 17 / Val Distance 4.685174418604651 / Train Loss 0.174586221575737\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 18 / Val Distance 4.542490310077519 / Train Loss 0.1761394590139389\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 19 / Val Distance 4.6486434108527135 / Train Loss 0.166859433054924\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 20 / Val Distance 4.599006782945737 / Train Loss 0.16148658096790314\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 21 / Val Distance 4.516763565891472 / Train Loss 0.15553073585033417\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 22 / Val Distance 4.52109980620155 / Train Loss 0.1541569083929062\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 23 / Val Distance 4.44031007751938 / Train Loss 0.1481846123933792\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 24 / Val Distance 4.414970930232558 / Train Loss 0.14436665177345276\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 25 / Val Distance 4.217926356589147 / Train Loss 0.1400662362575531\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 26 / Val Distance 4.205983527131783 / Train Loss 0.13401147723197937\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 27 / Val Distance 4.208938953488372 / Train Loss 0.13046687841415405\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 28 / Val Distance 4.163687015503876 / Train Loss 0.12831848859786987\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 29 / Val Distance 4.075121124031008 / Train Loss 0.12340879440307617\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 30 / Val Distance 4.021148255813954 / Train Loss 0.11604584753513336\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 31 / Val Distance 3.9989825581395353 / Train Loss 0.11167296767234802\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 32 / Val Distance 3.9367490310077518 / Train Loss 0.10831940919160843\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 33 / Val Distance 4.002737403100775 / Train Loss 0.10394959896802902\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 34 / Val Distance 3.9075096899224806 / Train Loss 0.09959633648395538\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 35 / Val Distance 3.968556201550388 / Train Loss 0.09550676494836807\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 36 / Val Distance 3.8169573643410852 / Train Loss 0.09162404388189316\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 37 / Val Distance 3.8815164728682165 / Train Loss 0.08679598569869995\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 38 / Val Distance 3.8546027131782945 / Train Loss 0.0826227217912674\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 39 / Val Distance 3.7175387596899223 / Train Loss 0.08275941759347916\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 40 / Val Distance 3.7438226744186047 / Train Loss 0.07315197587013245\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 41 / Val Distance 3.6726744186046516 / Train Loss 0.07174447923898697\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 42 / Val Distance 3.7292635658914723 / Train Loss 0.06793301552534103\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 43 / Val Distance 3.6758963178294577 / Train Loss 0.06499616801738739\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 44 / Val Distance 3.6088905038759687 / Train Loss 0.06192905455827713\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 45 / Val Distance 3.6247335271317827 / Train Loss 0.05740424990653992\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 46 / Val Distance 3.583987403100775 / Train Loss 0.05465194955468178\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 47 / Val Distance 3.6097868217054265 / Train Loss 0.05155617743730545\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 48 / Val Distance 3.591109496124031 / Train Loss 0.04983055591583252\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 49 / Val Distance 3.577543604651163 / Train Loss 0.04611481353640556\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 50 / Val Distance 3.5465843023255816 / Train Loss 0.042998407036066055\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 51 / Val Distance 3.49859496124031 / Train Loss 0.04173871502280235\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 52 / Val Distance 3.5266472868217056 / Train Loss 0.03835737332701683\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 53 / Val Distance 3.4695494186046516 / Train Loss 0.03559817001223564\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 54 / Val Distance 3.4412063953488374 / Train Loss 0.03455283120274544\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 55 / Val Distance 3.5104651162790694 / Train Loss 0.03228653967380524\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 56 / Val Distance 3.481516472868217 / Train Loss 0.030205491930246353\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 57 / Val Distance 3.4453972868217053 / Train Loss 0.02817699871957302\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 58 / Val Distance 3.356904069767442 / Train Loss 0.02659718133509159\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 59 / Val Distance 3.423013565891473 / Train Loss 0.025379138067364693\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 60 / Val Distance 3.388032945736434 / Train Loss 0.023119056597352028\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 61 / Val Distance 3.3992005813953488 / Train Loss 0.02144857682287693\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 62 / Val Distance 3.359205426356589 / Train Loss 0.020321402698755264\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 63 / Val Distance 3.4009205426356592 / Train Loss 0.018790267407894135\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 64 / Val Distance 3.342126937984496 / Train Loss 0.01729011908173561\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 65 / Val Distance 3.3481589147286823 / Train Loss 0.01641547866165638\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 66 / Val Distance 3.3854893410852713 / Train Loss 0.015320675447583199\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 67 / Val Distance 3.363493217054264 / Train Loss 0.014028869569301605\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 68 / Val Distance 3.332485465116279 / Train Loss 0.01328524574637413\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 69 / Val Distance 3.295251937984496 / Train Loss 0.012263513170182705\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 70 / Val Distance 3.3188953488372093 / Train Loss 0.011422162875533104\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 71 / Val Distance 3.297456395348837 / Train Loss 0.010625881142914295\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 72 / Val Distance 3.299539728682171 / Train Loss 0.009977847337722778\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 73 / Val Distance 3.293095930232558 / Train Loss 0.009133453480899334\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 74 / Val Distance 3.283187984496124 / Train Loss 0.008553392253816128\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 75 / Val Distance 3.2859496124031007 / Train Loss 0.007839558646082878\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 76 / Val Distance 3.30390019379845 / Train Loss 0.0072490256279706955\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 77 / Val Distance 3.2303052325581394 / Train Loss 0.006857526954263449\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 78 / Val Distance 3.268580426356589 / Train Loss 0.006376036908477545\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 79 / Val Distance 3.227592054263566 / Train Loss 0.005851452238857746\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 80 / Val Distance 3.211967054263566 / Train Loss 0.0054629179649055\n",
            "Saving model\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 81 / Val Distance 3.2480135658914726 / Train Loss 0.005189976654946804\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            ""
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 82 / Val Distance 3.23718507751938 / Train Loss 0.004766519647091627\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Train:  22%|██▏       | 365/1626 [07:37<26:26,  1.26s/it, lr=6.784478222989877e-05, tloss=0.0046]"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-34-e99f717ca526>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m     \u001b[0;31m# Where you have your scheduler.step depends on the scheduler you use.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m     \u001b[0mtrain_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mscheduler\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mscaler\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m     \u001b[0mval_dist\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mval_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mval_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"\\nEpoch {epoch} / Val Distance {val_dist} / Train Loss {train_loss}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-33-26b9e6b87eee>\u001b[0m in \u001b[0;36mtrain_step\u001b[0;34m(train_loader, model, optimizer, criterion, scheduler, scaler)\u001b[0m\n\u001b[1;32m     36\u001b[0m         \u001b[0mtrain_loss\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m         \u001b[0mscaler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 38\u001b[0;31m         \u001b[0mscaler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     39\u001b[0m         \u001b[0mscaler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m         \u001b[0mscheduler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/cuda/amp/grad_scaler.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, optimizer, *args, **kwargs)\u001b[0m\n\u001b[1;32m    336\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moptimizer_state\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"found_inf_per_device\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"No inf checks were recorded for this optimizer.\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    337\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 338\u001b[0;31m         \u001b[0mretval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maybe_opt_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    339\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    340\u001b[0m         \u001b[0moptimizer_state\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"stage\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mOptState\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSTEPPED\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/cuda/amp/grad_scaler.py\u001b[0m in \u001b[0;36m_maybe_opt_step\u001b[0;34m(self, optimizer, optimizer_state, *args, **kwargs)\u001b[0m\n\u001b[1;32m    282\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_maybe_opt_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    283\u001b[0m         \u001b[0mretval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 284\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mv\u001b[0m \u001b[0;32min\u001b[0m \u001b[0moptimizer_state\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"found_inf_per_device\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    285\u001b[0m             \u001b[0mretval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    286\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mretval\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/cuda/amp/grad_scaler.py\u001b[0m in \u001b[0;36m<genexpr>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    282\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_maybe_opt_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    283\u001b[0m         \u001b[0mretval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 284\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mv\u001b[0m \u001b[0;32min\u001b[0m \u001b[0moptimizer_state\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"found_inf_per_device\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    285\u001b[0m             \u001b[0mretval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    286\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mretval\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate Predictions and Submit to Kaggle"
      ],
      "metadata": {
        "id": "M2H4EEj-sD32"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "fvDQCGAYOdHB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "evaluate(val_loader, model)"
      ],
      "metadata": {
        "id": "KtlU8yRn24ys"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "evaluate(val_loader, model)"
      ],
      "metadata": {
        "id": "EH1uaha5Y13i",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cd341533-32af-48dc-bead-7d35073d47ae"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(3.2595687984496124, tensor(0.3999, device='cuda:0'))"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#TODO: Make predictions\n",
        "\n",
        "# Follow the steps below:\n",
        "# 1. Create a new object for CTCBeamDecoder with larger (why?) number of beams\n",
        "# 2. Get prediction string by decoding the results of the beam decoder\n",
        "\n",
        "decoder_test = CTCBeamDecoder(\n",
        "    ARPAbet,\n",
        "    model_path=None,\n",
        "    alpha=0,\n",
        "    beta=0,\n",
        "    cutoff_top_n=40,\n",
        "    cutoff_prob=1.0,\n",
        "    beam_width=train_config['beam_width'],\n",
        "    num_processes=4,\n",
        "    blank_id=0,\n",
        "    log_probs_input=True)\n",
        "\n",
        "def make_output(h, lh, decoder, LABELS):\n",
        "  \n",
        "    beam_results, beam_scores, timesteps, out_seq_len = decoder_test.decode(h, lh) #TODO: What parameters would the decode function take in?\n",
        "    batch_size = h.shape[0] #What is the batch size\n",
        "\n",
        "    preds = []\n",
        "    for i in range(batch_size): # Loop through each element in the batch\n",
        "\n",
        "        h_slice = beam_results[i, 0, :out_seq_len[i, 0]]\n",
        "        h_string = [LABELS[x] for x in h_slice]\n",
        "        preds.append(\"\".join(h_string))\n",
        "    \n",
        "    return preds"
      ],
      "metadata": {
        "id": "2moYJhTWsOG-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def predict(dataloader, model, decoder, LABELS):\n",
        "    model.eval()\n",
        "\n",
        "    test_results = []\n",
        "    with torch.no_grad():\n",
        "      for i, data in enumerate(dataloader):\n",
        "          h, lh = data\n",
        "          h = h.to(device)\n",
        "\n",
        "          # HINT: Are you using mixed precision?\n",
        "\n",
        "          with torch.cuda.amp.autocast():\n",
        "              out, out_len = model(h, lh)\n",
        "          pred = make_output(out, out_len, decoder_test, LABELS)\n",
        "          test_results.extend(pred)\n",
        "          del h, lh, out, out_len\n",
        "          torch.cuda.empty_cache()\n",
        "          gc.collect()\n",
        "    \n",
        "    return test_results"
      ],
      "metadata": {
        "id": "W2VCfRFQUhoy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#TODO:\n",
        "# Write a function (predict) to generate predictions and submit the file to Kaggle\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "predictions = predict(test_loader, model, decoder_test, LABELS)\n",
        "import pandas as pd\n",
        "\n",
        "df = pd.read_csv('/content/hw3p2/test-clean/transcript/random_submission.csv')\n",
        "df.label = predictions\n",
        "\n",
        "df.to_csv('submission.csv', index = False)\n",
        "!kaggle competitions submit -c '11-785-f22-hw3p2' -f submission.csv -m \"I made it!\""
      ],
      "metadata": {
        "id": "qESloOg5hpgZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9ddacd3c-613f-44c4-eda5-bf745fe92af5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Warning: Looks like you're using an outdated API Version, please consider updating (server 1.5.12 / client 1.5.8)\n",
            "100% 210k/210k [00:00<00:00, 965kB/s]\n",
            "Successfully submitted to Automatic Speech Recognition (ASR)"
          ]
        }
      ]
    }
  ]
}