examples/asr/notebooks/3_Speech_Commands_using_NeMo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "!pip install wget\n",
    "!pip install git+https://github.com/NVIDIA/apex.git\n",
    "!pip install nemo-toolkit\n",
    "!pip install nemo-asr\n",
    "!pip install unidecode\n",
    "\n",
    "!mkdir configs\n",
    "!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/master/examples/asr/configs/quartznet_speech_commands_3x1_v1.yaml\n",
    "!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/master/examples/asr/configs/quartznet_speech_commands_3x1_v2.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import some necessary libraries\n",
    "import os\n",
    "import argparse\n",
    "import copy\n",
    "import math\n",
    "import os\n",
    "import glob\n",
    "from functools import partial\n",
    "from datetime import datetime\n",
    "from ruamel.yaml import YAML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This Speech Command recognition tutorial is based on the QuartzNet model from the paper \"[QuartzNet: Deep Automatic Speech Recognition with 1D Time-Channel Separable Convolutions](https://arxiv.org/pdf/1910.10261.pdf)\" with a modified decoder head to suit classification tasks.\n",
    "\n",
    "The notebook will follow the steps below:\n",
    "\n",
    " - Dataset preparation: Preparing Google Speech Commands dataset\n",
    "\n",
    " - Audio preprocessing (feature extraction): signal normalization, windowing, (log) spectrogram (or mel scale spectrogram, or MFCC)\n",
    "\n",
    " - Data augmentation using SpecAugment \"[SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition](https://arxiv.org/abs/1904.08779)\" to increase number of data samples.\n",
    " \n",
    " - Develop a small Neural classification model which can be trained efficiently.\n",
    " \n",
    " - Model training on the Google Speech Commands dataset in NeMo.\n",
    " \n",
    " - Evaluation of error cases of the model by audibly hearing the samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is where the Google Speech Commands directory will be placed.\n",
    "# Change this if you don't want the data to be extracted in the current directory.\n",
    "# Select the version of the dataset required as well (can be 1 or 2)\n",
    "DATASET_VER = 1\n",
    "data_dir = './google_dataset_v{0}/'.format(DATASET_VER)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Preparation\n",
    "\n",
    "We will be using the open source Google Speech Commands Dataset (we will use V1 of the dataset for the tutorial, but require very minor changes to support V2 dataset). These scripts below will download the dataset and convert it to a format suitable for use with nemo_asr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the dataset\n",
    "\n",
    "The dataset must be prepared using the scripts provided under the `{NeMo root directory}/scripts` sub-directory. \n",
    "\n",
    "Run the following command below to download the training script and execute it.\n",
    "\n",
    "**NOTE**: You should have at least 4GB of disk space available if you’ve used --data_version=1; and at least 6GB if you used --data_version=2. Also, it will take some time to download and process, so go grab a coffee.\n",
    "\n",
    "**NOTE**: You may additionally pass a `--rebalance` flag at the end of the `process_speech_commands_data.py` script to rebalance the class samples in the manifest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2020-02-27 16:14:08--  https://raw.githubusercontent.com/NVIDIA/NeMo/master/scripts/process_speech_commands_data.py\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 151.101.40.133\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|151.101.40.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 6872 (6.7K) [text/plain]\n",
      "Saving to: ‘process_speech_commands_data.py.1’\n",
      "\n",
      "process_speech_comm 100%[===================>]   6.71K  --.-KB/s    in 0s      \n",
      "\n",
      "2020-02-27 16:14:08 (58.8 MB/s) - ‘process_speech_commands_data.py.1’ saved [6872/6872]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://raw.githubusercontent.com/NVIDIA/NeMo/master/scripts/process_speech_commands_data.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘./google_dataset_v1/’: File exists\n",
      "Dataset ready !\n"
     ]
    }
   ],
   "source": [
    "!mkdir {data_dir}\n",
    "!python process_speech_commands_data.py --data_root={data_dir} --data_version={DATASET_VER}\n",
    "print(\"Dataset ready !\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the path to manifest files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtaset_path = 'google_speech_recognition_v{0}'.format(DATASET_VER)\n",
    "dataset_basedir = os.path.join(data_dir, dtaset_path)\n",
    "\n",
    "train_dataset = os.path.join(dataset_basedir, 'train_manifest.json')\n",
    "val_dataset = os.path.join(dataset_basedir, 'validation_manifest.json')\n",
    "test_dataset = os.path.join(dataset_basedir, 'validation_manifest.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read a few rows of the manifest file \n",
    "\n",
    "Manifest files are the data structure used by NeMo to declare a few important details about the data :\n",
    "\n",
    "1) `audio_filepath`: Refers to the path to the raw audio file <br>\n",
    "2) `command`: The class label (or speech command) of this sample <br>\n",
    "3) `duration`: The length of the audio file, in seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "head: cannot open '{train_manifest}' for reading: No such file or directory\n"
     ]
    }
   ],
   "source": [
    "!head -n 5 {train_dataset}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training - Preparation\n",
    "\n",
    "We will be training a QuartzNet model from the paper \"[QuartzNet: Deep Automatic Speech Recognition with 1D Time-Channel Separable Convolutions](https://arxiv.org/pdf/1910.10261.pdf)\". The benefit of QuartzNet over JASPER models is that they use Separable Convolutions, which greatly reduce the number of parameters required to get good model accuracy.\n",
    "\n",
    "QuartzNet models generally follow the model definition pattern QuartzNet-[BxR], where B is the number of blocks and R is the number of convolutional sub-blocks. Each sub-block contains a 1-D masked convolution, batch normalization, ReLU, and dropout:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets load the config file for the QuartzNet 3x1 model\n",
    "# Here we will be using separable convolutions\n",
    "# with 3 blocks (k=3 repeated once r=1 from the picture above)\n",
    "yaml = YAML(typ=\"safe\")\n",
    "with open(\"configs/quartznet_speech_commands_3x1_v{0}.yaml\".format(DATASET_VER)) as f:\n",
    "    jasper_params = yaml.load(f)\n",
    "\n",
    "# Pre-define a set of labels that this model must learn to predict\n",
    "labels = jasper_params['labels']\n",
    "\n",
    "# Get the sampling rate of the data\n",
    "sample_rate = jasper_params['sample_rate']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import NeMo core functionality\n",
    "# NeMo's \"core\" package\n",
    "import nemo\n",
    "# NeMo's ASR collection\n",
    "import nemo.collections.asr as nemo_asr\n",
    "# NeMo's learning rate policy\n",
    "from nemo.utils.lr_policies import CosineAnnealing\n",
    "from nemo.collections.asr.helpers import (\n",
    "    monitor_classification_training_progress,\n",
    "    process_classification_evaluation_batch,\n",
    "    process_classification_evaluation_epoch,\n",
    ")\n",
    "from nemo.collections.asr.metrics import classification_accuracy\n",
    "\n",
    "logging = nemo.logging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define some model hyper parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets define some hyper parameters\n",
    "lr = 0.05\n",
    "num_epochs = 5\n",
    "batch_size = 128\n",
    "weight_decay = 0.001"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the NeMo components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2020-02-27 16:14:20 deprecated:68] Function ``_get_trainer`` is deprecated. It is going to be removed in the future version.\n"
     ]
    }
   ],
   "source": [
    "# Create a Neural Factory\n",
    "# It creates log files and tensorboard writers for us among other functions\n",
    "neural_factory = nemo.core.NeuralModuleFactory(\n",
    "    log_dir='./{0}/quartznet-3x1-v{1}'.format(dataset_basedir, DATASET_VER),\n",
    "    create_tb_writer=True)\n",
    "tb_writer = neural_factory.tb_writer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:21 collections:215] Filtered duration for loading collection is 0.000000.\n",
      "[NeMo I 2020-02-27 16:14:21 collections:215] Filtered duration for loading collection is 0.000000.\n",
      "[NeMo I 2020-02-27 16:14:21 features:144] PADDING: 16\n",
      "[NeMo I 2020-02-27 16:14:21 features:152] STFT using conv\n",
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-12-242cb97ccf7d>:34] Steps per epoch : 401\n",
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-12-242cb97ccf7d>:35] Have 51088 examples to train on.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='mean' instead.\n",
      "  warnings.warn(warning.format(ret))\n"
     ]
    }
   ],
   "source": [
    "# Check if data augmentation such as white noise and time shift augmentation should be used\n",
    "audio_augmentor = jasper_params.get('AudioAugmentor', None)\n",
    "\n",
    "# Build the input data layer and the preprocessing layers for the train set\n",
    "train_data_layer = nemo_asr.AudioToSpeechLabelDataLayer(\n",
    "    manifest_filepath=train_dataset,\n",
    "    labels=labels,\n",
    "    sample_rate=sample_rate,\n",
    "    batch_size=batch_size,\n",
    "    num_workers=os.cpu_count(),\n",
    "    augmentor=audio_augmentor,\n",
    "    shuffle=True\n",
    ")\n",
    "\n",
    " # Build the input data layer and the preprocessing layers for the test set\n",
    "eval_data_layer = nemo_asr.AudioToSpeechLabelDataLayer(\n",
    "    manifest_filepath=test_dataset,\n",
    "    sample_rate=sample_rate,\n",
    "    labels=labels,\n",
    "    batch_size=batch_size,\n",
    "    num_workers=os.cpu_count(),\n",
    "    shuffle=False,\n",
    ")\n",
    "\n",
    "# We will convert the raw audio data into MelSpectrogram Features to feed as input to our model\n",
    "data_preprocessor = nemo_asr.AudioToMelSpectrogramPreprocessor(\n",
    "    sample_rate=sample_rate, **jasper_params[\"AudioToMelSpectrogramPreprocessor\"],\n",
    ")\n",
    "\n",
    "# Compute the total number of samples and the number of training steps per epoch\n",
    "N = len(train_data_layer)\n",
    "steps_per_epoch = math.ceil(N / float(batch_size) + 1)\n",
    "\n",
    "logging.info(\"Steps per epoch : {0}\".format(steps_per_epoch))\n",
    "logging.info('Have {0} examples to train on.'.format(N))\n",
    "\n",
    "# Here we begin defining all of the augmentations we want\n",
    "# We will pad the preprocessed spectrogram image to have a certain number of timesteps\n",
    "# This centers the generated spectrogram and adds black boundaries to either side\n",
    "# of the padded image.\n",
    "crop_pad_augmentation = nemo_asr.CropOrPadSpectrogramAugmentation(audio_length=128)\n",
    "\n",
    "# We also optionally add `SpecAugment` augmentations based on the config file\n",
    "# SpecAugment has various possible augmentations to the generated spectrogram\n",
    "# 1) Frequency band masking\n",
    "# 2) Time band masking\n",
    "# 3) Rectangular cutout\n",
    "spectr_augment_config = jasper_params.get('SpectrogramAugmentation', None)\n",
    "if spectr_augment_config:\n",
    "    data_spectr_augmentation = nemo_asr.SpectrogramAugmentation(**spectr_augment_config)\n",
    "\n",
    "# Build the QuartzNet Encoder model\n",
    "# The config defines the layers as a list of dictionaries\n",
    "# The first and last two blocks are not considered when we say QuartzNet-[BxR]\n",
    "# B is counted as the number of blocks after the first layer and before the penultimate layer.\n",
    "# R is defined as the number of repetitions of each block in B.\n",
    "# Note: We can scale the convolution kernels size by the float parameter `kernel_size_factor`\n",
    "jasper_encoder = nemo_asr.JasperEncoder(**jasper_params[\"JasperEncoder\"])\n",
    "\n",
    "# We then define the QuartzNet decoder.\n",
    "# This decoder head is specialized for the task for classification, such that it\n",
    "# accepts a set of `N-feat` per timestep of the model, and averages these features\n",
    "# over all the timesteps, before passing a Linear classification layer on those features.\n",
    "jasper_decoder = nemo_asr.JasperDecoderForClassification(\n",
    "    feat_in=jasper_params[\"JasperEncoder\"][\"jasper\"][-1][\"filters\"],\n",
    "    num_classes=len(labels),\n",
    "    **jasper_params['JasperDecoderForClassification'],\n",
    ")\n",
    "\n",
    "# We can easily apply cross entropy loss to train this model\n",
    "ce_loss = nemo_asr.CrossEntropyLossNM()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-13-6805b5462cf6>:2] ================================\n",
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-13-6805b5462cf6>:3] Number of parameters in encoder: 73344\n",
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-13-6805b5462cf6>:4] Number of parameters in decoder: 3870\n",
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-13-6805b5462cf6>:6] Total number of parameters in model: 77214\n",
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-13-6805b5462cf6>:8] ================================\n"
     ]
    }
   ],
   "source": [
    "# Lets print out the number of parameters of this model\n",
    "logging.info('================================')\n",
    "logging.info(f\"Number of parameters in encoder: {jasper_encoder.num_weights}\")\n",
    "logging.info(f\"Number of parameters in decoder: {jasper_decoder.num_weights}\")\n",
    "logging.info(\n",
    "    f\"Total number of parameters in model: \" f\"{jasper_decoder.num_weights + jasper_encoder.num_weights}\"\n",
    ")\n",
    "logging.info('================================')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile the Training Graph for NeMo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we have all of the components that are required to build the NeMo execution graph!\n",
    "## Build the training data loaders and preprocessors first\n",
    "audio_signal, audio_signal_len, commands, command_len = train_data_layer()\n",
    "processed_signal, processed_signal_len = data_preprocessor(input_signal=audio_signal, length=audio_signal_len)\n",
    "processed_signal, processed_signal_len = crop_pad_augmentation(\n",
    "    input_signal=processed_signal,\n",
    "    length=audio_signal_len\n",
    ")\n",
    "\n",
    "## Augment the dataset for training\n",
    "if spectr_augment_config:\n",
    "    processed_signal = data_spectr_augmentation(input_spec=processed_signal)\n",
    "\n",
    "## Define the model\n",
    "encoded, encoded_len = jasper_encoder(audio_signal=processed_signal, length=processed_signal_len)\n",
    "decoded = jasper_decoder(encoder_output=encoded)\n",
    "\n",
    "## Obtain the train loss\n",
    "train_loss = ce_loss(logits=decoded, labels=commands)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile the Test Graph for NeMo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we build the test graph in a similar way, reusing the above components\n",
    "## Build the test data loader and preprocess same way as train graph\n",
    "## But note, we do not add the spectrogram augmentation to the test graph !\n",
    "test_audio_signal, test_audio_signal_len, test_commands, test_command_len = eval_data_layer()\n",
    "test_processed_signal, test_processed_signal_len = data_preprocessor(\n",
    "    input_signal=test_audio_signal, length=test_audio_signal_len\n",
    ")\n",
    "test_processed_signal, test_processed_signal_len = crop_pad_augmentation(\n",
    "    input_signal=test_processed_signal, length=test_processed_signal_len\n",
    ")\n",
    "\n",
    "# Pass the test data through the model encoder and decoder\n",
    "test_encoded, test_encoded_len = jasper_encoder(\n",
    "    audio_signal=test_processed_signal, length=test_processed_signal_len\n",
    ")\n",
    "test_decoded = jasper_decoder(encoder_output=test_encoded)\n",
    "\n",
    "# Compute test loss for visualization\n",
    "test_loss = ce_loss(logits=test_decoded, labels=test_commands)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up callbacks for training and test set evaluation, and checkpoint saving"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now that we have our training and evaluation graphs built,\n",
    "# we can focus on a few callbacks to help us save the model checkpoints\n",
    "# during training, as well as display train and test metrics\n",
    "\n",
    "# Callbacks needed to print train info to console and Tensorboard\n",
    "train_callback = nemo.core.SimpleLossLoggerCallback(\n",
    "    # Notice that we pass in loss, predictions, and the labels.\n",
    "    # Of course we would like to see our training loss, but we need the\n",
    "    # other arguments to calculate the accuracy.\n",
    "    tensors=[train_loss, decoded, commands],\n",
    "    # The print_func defines what gets printed.\n",
    "    print_func=partial(monitor_classification_training_progress, eval_metric=None),\n",
    "    get_tb_values=lambda x: [(\"loss\", x[0])],\n",
    "    tb_writer=neural_factory.tb_writer,\n",
    ")\n",
    "\n",
    "# Callbacks needed to print test info to console and Tensorboard\n",
    "tagname = 'TestSet'\n",
    "eval_callback = nemo.core.EvaluatorCallback(\n",
    "    eval_tensors=[test_loss, test_decoded, test_commands],\n",
    "    user_iter_callback=partial(process_classification_evaluation_batch, top_k=1),\n",
    "    user_epochs_done_callback=partial(process_classification_evaluation_epoch, eval_metric=1, tag=tagname),\n",
    "    eval_step=200,  # How often we evaluate the model on the test set\n",
    "    tb_writer=neural_factory.tb_writer,\n",
    ")\n",
    "\n",
    "# Callback to save model checkpoints\n",
    "chpt_callback = nemo.core.CheckpointCallback(\n",
    "    folder=neural_factory.checkpoint_dir,\n",
    "    step_freq=1000,\n",
    ")\n",
    "\n",
    "# Prepare a list of checkpoints to pass to the engine\n",
    "callbacks = [train_callback, eval_callback, chpt_callback]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training the model\n",
    "\n",
    "Even with such a small model (77k parameters), and just 5 epochs (should take just a few minutes to train), you should be able to get a test set accuracy score in the range 85 - 90%. Not bad for a 30 (v1) or 35 (v2) way classification problem !\n",
    "\n",
    "Experiment with increasing the number of epochs or with batch size to see how much you can improve the score!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:24 <ipython-input-17-e865cc4031ec>:11] Using `<nemo.utils.lr_policies.CosineAnnealing object at 0x7f547c65e210>` Learning Rate Scheduler\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:179] Starting .....\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:343] Found 2 modules with weights:\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:345] JasperEncoder\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:345] JasperDecoderForClassification\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:346] Total model parameters: 77214\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:301] Restoring checkpoint from folder ././google_dataset_v1/google_speech_recognition_v1/quartznet-3x1-v1/checkpoints ...\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:186] Done in 0.025618553161621094\n",
      "[NeMo I 2020-02-27 16:14:24 callbacks:432] Final Evaluation ..............................\n",
      "[NeMo I 2020-02-27 16:14:27 helpers:273] ==========>>>>>>Evaluation Loss TestSet: 0.35130253434181213\n",
      "[NeMo I 2020-02-27 16:14:27 helpers:275] ==========>>>>>>Evaluation Accuracy Top@1 TestSet: 89.5999\n",
      "[NeMo I 2020-02-27 16:14:27 callbacks:437] Evaluation time: 2.7541420459747314 seconds\n",
      "[NeMo I 2020-02-27 16:14:27 callbacks:293] Saved checkpoint: ././google_dataset_v1/google_speech_recognition_v1/quartznet-3x1-v1/checkpoints/trainer-STEP-2000.pt\n"
     ]
    }
   ],
   "source": [
    "# Now we have all the components required to train the model\n",
    "# Lets define a learning rate schedule\n",
    "\n",
    "# Define a learning rate schedule\n",
    "lr_policy = CosineAnnealing(\n",
    "    total_steps=num_epochs * steps_per_epoch,\n",
    "    warmup_ratio=0.05,\n",
    "    min_lr=0.001,\n",
    ")\n",
    "\n",
    "logging.info(f\"Using `{lr_policy}` Learning Rate Scheduler\")\n",
    "\n",
    "# Finally, lets train this model !\n",
    "neural_factory.train(\n",
    "    tensors_to_optimize=[train_loss],\n",
    "    callbacks=callbacks,\n",
    "    lr_policy=lr_policy,\n",
    "    optimizer=\"novograd\",\n",
    "    optimization_params={\n",
    "        \"num_epochs\": num_epochs,\n",
    "        \"max_steps\": None,\n",
    "        \"lr\": lr,\n",
    "        \"momentum\": 0.95,\n",
    "        \"betas\": (0.98, 0.5),\n",
    "        \"weight_decay\": weight_decay,\n",
    "        \"grad_norm_clip\": None,\n",
    "    },\n",
    "    batches_per_step=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation of incorrectly predicted samples\n",
    "\n",
    "Given that we have a trained model, which performs reasonably well, lets try to listen to the samples where the model is least confident in its predictions.\n",
    "\n",
    "For this, we need support of the librosa library.\n",
    "\n",
    "**NOTE**: The following code depends on librosa. To install it, run the following code block first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: librosa in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (0.7.2)\n",
      "Requirement already satisfied: decorator>=3.0.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (4.4.1)\n",
      "Requirement already satisfied: resampy>=0.2.2 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (0.2.2)\n",
      "Requirement already satisfied: joblib>=0.12 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (0.14.1)\n",
      "Requirement already satisfied: scipy>=1.0.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (1.4.1)\n",
      "Requirement already satisfied: numba>=0.43.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (0.47.0)\n",
      "Requirement already satisfied: scikit-learn!=0.19.0,>=0.14.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (0.22.1)\n",
      "Requirement already satisfied: six>=1.3 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (1.14.0)\n",
      "Requirement already satisfied: numpy>=1.15.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (1.18.1)\n",
      "Requirement already satisfied: audioread>=2.0.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (2.1.8)\n",
      "Requirement already satisfied: soundfile>=0.9.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from librosa) (0.10.3.post1)\n",
      "Requirement already satisfied: setuptools in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from numba>=0.43.0->librosa) (44.0.0.post20200106)\n",
      "Requirement already satisfied: llvmlite>=0.31.0dev0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from numba>=0.43.0->librosa) (0.31.0)\n",
      "Requirement already satisfied: cffi>=1.0 in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from soundfile>=0.9.0->librosa) (1.13.2)\n",
      "Requirement already satisfied: pycparser in /home/smajumdar/anaconda3/envs/NeMo/lib/python3.7/site-packages (from cffi>=1.0->soundfile>=0.9.0->librosa) (2.19)\n"
     ]
    }
   ],
   "source": [
    "!pip install librosa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets add a path to the checkpoint dir\n",
    "model_path = neural_factory.checkpoint_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract the predictions from the model\n",
    "\n",
    "We want to possess the actual logits of the model instead of just the final evaluation score, so we use `NeuralFactory.infer(...)` to extract the logits per batch of samples provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:28 actions:1453] Restoring JasperEncoder from ././google_dataset_v1/google_speech_recognition_v1/quartznet-3x1-v1/checkpoints/JasperEncoder-STEP-2000.pt\n",
      "[NeMo I 2020-02-27 16:14:28 actions:1453] Restoring JasperDecoderForClassification from ././google_dataset_v1/google_speech_recognition_v1/quartznet-3x1-v1/checkpoints/JasperDecoderForClassification-STEP-2000.pt\n",
      "[NeMo I 2020-02-27 16:14:29 actions:726] Evaluating batch 0 out of 54\n",
      "[NeMo I 2020-02-27 16:14:29 actions:726] Evaluating batch 5 out of 54\n",
      "[NeMo I 2020-02-27 16:14:29 actions:726] Evaluating batch 10 out of 54\n",
      "[NeMo I 2020-02-27 16:14:29 actions:726] Evaluating batch 15 out of 54\n",
      "[NeMo I 2020-02-27 16:14:30 actions:726] Evaluating batch 20 out of 54\n",
      "[NeMo I 2020-02-27 16:14:30 actions:726] Evaluating batch 25 out of 54\n",
      "[NeMo I 2020-02-27 16:14:30 actions:726] Evaluating batch 30 out of 54\n",
      "[NeMo I 2020-02-27 16:14:30 actions:726] Evaluating batch 35 out of 54\n",
      "[NeMo I 2020-02-27 16:14:30 actions:726] Evaluating batch 40 out of 54\n",
      "[NeMo I 2020-02-27 16:14:31 actions:726] Evaluating batch 45 out of 54\n",
      "[NeMo I 2020-02-27 16:14:31 actions:726] Evaluating batch 50 out of 54\n"
     ]
    }
   ],
   "source": [
    "# --- Inference Only --- #\n",
    "# We've already built the inference DAG above, so all we need is to call infer().\n",
    "evaluated_tensors = neural_factory.infer(\n",
    "    # These are the tensors we want to get from the model.\n",
    "    tensors=[test_loss, test_decoded, test_commands],\n",
    "    # checkpoint_dir specifies where the model params are loaded from.\n",
    "    checkpoint_dir=model_path\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accuracy calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-21-674fb7de9132>:19] Total correct / Total count : 6094 / 6798\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-21-674fb7de9132>:20] Final accuracy : 0.8964401294498382\n"
     ]
    }
   ],
   "source": [
    "correct_count = 0\n",
    "total_count = 0\n",
    "\n",
    "for batch_idx, (logits, labels) in enumerate(zip(evaluated_tensors[1], evaluated_tensors[2])):\n",
    "    acc = classification_accuracy(\n",
    "        logits=logits,\n",
    "        targets=labels,\n",
    "        top_k=[1]\n",
    "    )\n",
    "\n",
    "    # Select top 1 accuracy only\n",
    "    acc = acc[0]\n",
    "\n",
    "    # Since accuracy here is \"per batch\", we simply denormalize it by multiplying\n",
    "    # by batch size to recover the count of correct samples.\n",
    "    correct_count += int(acc * logits.size(0))\n",
    "    total_count += logits.size(0)\n",
    "\n",
    "logging.info(f\"Total correct / Total count : {correct_count} / {total_count}\")\n",
    "logging.info(f\"Final accuracy : {correct_count / float(total_count)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filtering out incorrect samples\n",
    "Let us now filter out the incorrectly labeled samples from the total set of samples in the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import librosa\n",
    "import json\n",
    "import IPython.display as ipd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First lets create a utility class to remap the integer class labels to actual string label\n",
    "class ReverseMapLabel:\n",
    "    def __init__(self, data_layer: nemo_asr.AudioToSpeechLabelDataLayer):\n",
    "        self.label2id = dict(data_layer._dataset.label2id)\n",
    "        self.id2label = dict(data_layer._dataset.id2label)\n",
    "\n",
    "    def __call__(self, pred_idx, label_idx):\n",
    "        return self.id2label[pred_idx], self.id2label[label_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-24-3ed571e8b863>:22] Num test samples : 6798\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-24-3ed571e8b863>:23] Num errors : 704\n"
     ]
    }
   ],
   "source": [
    "# Next, lets get the indices of all the incorrectly labeled samples\n",
    "sample_idx = 0\n",
    "incorrect_preds = []\n",
    "rev_map = ReverseMapLabel(eval_data_layer)\n",
    "\n",
    "# Remember, evaluated_tensor = (loss, logits, labels)\n",
    "for batch_idx, (logits, labels) in enumerate(zip(evaluated_tensors[1], evaluated_tensors[2])):\n",
    "    probs = torch.softmax(logits, dim=-1)\n",
    "    probas, preds = torch.max(probs, dim=-1)\n",
    "\n",
    "    incorrect_ids = (preds != labels).nonzero()\n",
    "    for idx in incorrect_ids:\n",
    "        proba = float(probas[idx][0])\n",
    "        pred = int(preds[idx][0])\n",
    "        label = int(labels[idx][0])\n",
    "        idx = int(idx[0]) + sample_idx\n",
    "\n",
    "        incorrect_preds.append((idx, *rev_map(pred, label), proba))\n",
    "\n",
    "    sample_idx += labels.size(0)\n",
    "\n",
    "logging.info(f\"Num test samples : {total_count}\")\n",
    "logging.info(f\"Num errors : {len(incorrect_preds)}\")\n",
    "\n",
    "# First lets sort by confidence of prediction\n",
    "incorrect_preds = sorted(incorrect_preds, key=lambda x: x[-1], reverse=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examine a subset of incorrect samples\n",
    "Lets print out the (test id, predicted label, ground truth label, confidence) tuple of first 20 incorrectly labeled samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (3184, 'up', 'two', 0.13125509023666382)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (1966, 'wow', 'no', 0.13236339390277863)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (1415, 'up', 'yes', 0.13250434398651123)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (5428, 'nine', 'up', 0.13804833590984344)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (1837, 'up', 'zero', 0.1411990523338318)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (3083, 'four', 'two', 0.14131611585617065)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (885, 'one', 'eight', 0.143906369805336)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (5584, 'go', 'cat', 0.14928434789180756)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (6056, 'dog', 'sheila', 0.1584177315235138)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (5238, 'up', 'on', 0.15925118327140808)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (3401, 'up', 'go', 0.16477465629577637)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (4700, 'zero', 'off', 0.16728630661964417)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (5175, 'up', 'on', 0.1677844226360321)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (1353, 'three', 'yes', 0.17652447521686554)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (6284, 'nine', 'left', 0.18027563393115997)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (6123, 'nine', 'left', 0.18055355548858643)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (5793, 'two', 'stop', 0.18353451788425446)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (1136, 'one', 'nine', 0.18678408861160278)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (5042, 'bird', 'on', 0.1876409649848938)\n",
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-25-631305d430a9>:2] (807, 'up', 'eight', 0.19111238420009613)\n"
     ]
    }
   ],
   "source": [
    "for incorrect_sample in incorrect_preds[:20]:\n",
    "    logging.info(str(incorrect_sample))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Define a threshold below which we designate a model's prediction as \"low confidence\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2020-02-27 16:14:31 <ipython-input-26-a1b4199a519e>:4] Number of low confidence predictions : 39\n"
     ]
    }
   ],
   "source": [
    "# Filter out how many such samples exist\n",
    "low_confidence_threshold = 0.25\n",
    "count_low_confidence = len(list(filter(lambda x: x[-1] <= low_confidence_threshold, incorrect_preds)))\n",
    "logging.info(f\"Number of low confidence predictions : {count_low_confidence}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lets hear the samples which the model has least confidence in !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First lets create a helper function to parse the manifest files\n",
    "def parse_manifest(manifest):\n",
    "    data = []\n",
    "    for line in manifest:\n",
    "        line = json.loads(line)\n",
    "        data.append(line)\n",
    "\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, lets create a helper function to actually listen to certain samples\n",
    "def listen_to_file(sample_id, pred=None, label=None, proba=None):\n",
    "    # Load the audio waveform using librosa\n",
    "    filepath = test_samples[sample_id]['audio_filepath']\n",
    "    audio, sample_rate = librosa.load(filepath)\n",
    "\n",
    "    if pred is not None and label is not None and proba is not None:\n",
    "        logging.info(f\"Sample : {sample_id} Prediction : {pred} Label : {label} Confidence = {proba: 0.4f}\")\n",
    "    else:\n",
    "        logging.info(f\"Sample : {sample_id}\")\n",
    "\n",
    "    return ipd.Audio(audio, rate=sample_rate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now lets load the test manifest into memory\n",
    "test_samples = []\n",
    "with open(test_dataset, 'r') as test_f:\n",
    "    test_samples = test_f.readlines()\n",
    "\n",
    "test_samples = parse_manifest(test_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Finally, lets listen to all the audio samples where the model made a mistake\n",
    "# Note: This list of incorrect samples may be quite large, so you may choose to subsample `incorrect_preds`\n",
    "for sample_id, pred, label, proba in incorrect_preds[:count_low_confidence]:\n",
    "    ipd.display(listen_to_file(sample_id, pred=pred, label=label, proba=proba))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}