artcraft/models/tts/NeMo/tutorials/asr/03_Speech_Commands.ipynb

In [ ]:

"""
You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.

Instructions for setting up Colab are as follows:
1. Open a new Python 3 notebook.
2. Import this notebook from GitHub (File -> Upload Notebook -> "GITHUB" tab -> copy/paste GitHub URL)
3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select "GPU" for hardware accelerator)
4. Run this cell to set up dependencies.
"""
# If you're using Google Colab and not running locally, run this cell.

## Install dependencies
!pip install wget
!apt-get install sox libsndfile1 ffmpeg
!pip install unidecode

# ## Install NeMo
BRANCH = 'v1.0.2'
!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[asr]

## Install TorchAudio
!pip install torchaudio>=0.6.0 -f https://download.pytorch.org/whl/torch_stable.html

## Grab the config we'll use in this example
!mkdir configs

Introduction

This Speech Command recognition tutorial is based on the MatchboxNet model from the paper "MatchboxNet: 1D Time-Channel Separable Convolutional Neural Network Architecture for Speech Commands Recognition". MatchboxNet is a modified form of the QuartzNet architecture from the paper "QuartzNet: Deep Automatic Speech Recognition with 1D Time-Channel Separable Convolutions" with a modified decoder head to suit classification tasks.

The notebook will follow the steps below:

• Dataset preparation: Preparing Google Speech Commands dataset

• Audio preprocessing (feature extraction): signal normalization, windowing, (log) spectrogram (or mel scale spectrogram, or MFCC)

• Data augmentation using SpecAugment "SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition" to increase the number of data samples.

• Develop a small Neural classification model that can be trained efficiently.

• Model training on the Google Speech Commands dataset in NeMo.

• Evaluation of error cases of the model by audibly hearing the samples

In [ ]:

# Some utility imports
import os
from omegaconf import OmegaConf

In [ ]:

# This is where the Google Speech Commands directory will be placed.
# Change this if you don't want the data to be extracted in the current directory.
# Select the version of the dataset required as well (can be 1 or 2)
DATASET_VER = 1
data_dir = './google_dataset_v{0}/'.format(DATASET_VER)

if DATASET_VER == 1:
  MODEL_CONFIG = "matchboxnet_3x1x64_v1.yaml"
else:
  MODEL_CONFIG = "matchboxnet_3x1x64_v2.yaml"

if not os.path.exists(f"configs/{MODEL_CONFIG}"):
  !wget -P configs/ "https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/matchboxnet/{MODEL_CONFIG}"

Data Preparation

We will be using the open-source Google Speech Commands Dataset (we will use V1 of the dataset for the tutorial but require minor changes to support the V2 dataset). These scripts below will download the dataset and convert it to a format suitable for use with NeMo.

Download the dataset

The dataset must be prepared using the scripts provided under the {NeMo root directory}/scripts sub-directory.

Run the following command below to download the data preparation script and execute it.

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.

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.

In [ ]:

if not os.path.exists("process_speech_commands_data.py"):
  !wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/dataset_processing/process_speech_commands_data.py

Preparing the manifest file

The manifest file is a simple file that has the full path to the audio file, the duration of the audio file, and the label that is assigned to that audio file.

This notebook is only a demonstration, and therefore we will use the --skip_duration flag to speed up construction of the manifest file.

NOTE: When replicating the results of the paper, do not use this flag and prepare the manifest file with correct durations.

In [ ]:

!mkdir {data_dir}
!python process_speech_commands_data.py --data_root={data_dir} --data_version={DATASET_VER} --skip_duration --log
print("Dataset ready !")

Prepare the path to manifest files

In [ ]:

dataset_path = 'google_speech_recognition_v{0}'.format(DATASET_VER)
dataset_basedir = os.path.join(data_dir, dataset_path)

train_dataset = os.path.join(dataset_basedir, 'train_manifest.json')
val_dataset = os.path.join(dataset_basedir, 'validation_manifest.json')
test_dataset = os.path.join(dataset_basedir, 'validation_manifest.json')

Read a few rows of the manifest file

Manifest files are the data structure used by NeMo to declare a few important details about the data :

1. audio_filepath: Refers to the path to the raw audio file
   
2. command: The class label (or speech command) of this sample
   
3. duration: The length of the audio file, in seconds.

In [ ]:

!head -n 5 {train_dataset}

Training - Preparation

We will be training a MatchboxNet model from the paper "MatchboxNet: 1D Time-Channel Separable Convolutional Neural Network Architecture for Speech Commands Recognition". The benefit of MatchboxNet over JASPER models is that they use 1D Time-Channel Separable Convolutions, which greatly reduce the number of parameters required to obtain good model accuracy.

MatchboxNet models generally follow the model definition pattern QuartzNet-[BxRXC], where B is the number of blocks, R is the number of convolutional sub-blocks, and C is the number of channels in these blocks. Each sub-block contains a 1-D masked convolution, batch normalization, ReLU, and dropout.

An image of QuartzNet, the base configuration of MatchboxNet models, is provided below.

In [ ]:

# NeMo's "core" package
import nemo
# NeMo's ASR collection - this collections contains complete ASR models and
# building blocks (modules) for ASR
import nemo.collections.asr as nemo_asr

Model Configuration

The MatchboxNet Model is defined in a config file which declares multiple important sections.

They are:

1. model: All arguments that will relate to the Model - preprocessors, encoder, decoder, optimizer and schedulers, datasets and any other related information

2. trainer: Any argument to be passed to PyTorch Lightning

In [ ]:

# This line will print the entire config of the MatchboxNet model
config_path = f"configs/{MODEL_CONFIG}"
config = OmegaConf.load(config_path)
config = OmegaConf.to_container(config, resolve=True)
config = OmegaConf.create(config)
print(OmegaConf.to_yaml(config))

In [ ]:

# Preserve some useful parameters
labels = config.model.labels
sample_rate = config.sample_rate

Setting up the datasets within the config

If you'll notice, there are a few config dictionaries called train_ds, validation_ds and test_ds. These are configurations used to setup the Dataset and DataLoaders of the corresponding config.

In [ ]:

print(OmegaConf.to_yaml(config.model.train_ds))

??? inside configs

You will often notice that some configs have ??? in place of paths. This is used as a placeholder so that the user can change the value at a later time.

Let's add the paths to the manifests to the config above.

In [ ]:

config.model.train_ds.manifest_filepath = train_dataset
config.model.validation_ds.manifest_filepath = val_dataset
config.model.test_ds.manifest_filepath = test_dataset

Building the PyTorch Lightning Trainer

NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem!

Lets first instantiate a Trainer object!

In [ ]:

import torch
import pytorch_lightning as pl

In [ ]:

print("Trainer config - \n")
print(OmegaConf.to_yaml(config.trainer))

In [ ]:

# Lets modify some trainer configs for this demo
# Checks if we have GPU available and uses it
cuda = 1 if torch.cuda.is_available() else 0
config.trainer.gpus = cuda

# Reduces maximum number of epochs to 5 for quick demonstration
config.trainer.max_epochs = 5

# Remove distributed training flags
config.trainer.accelerator = None

In [ ]:

trainer = pl.Trainer(**config.trainer)

Setting up a NeMo Experiment

NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it !

In [ ]:

from nemo.utils.exp_manager import exp_manager

In [ ]:

exp_dir = exp_manager(trainer, config.get("exp_manager", None))

In [ ]:

# The exp_dir provides a path to the current experiment for easy access
exp_dir = str(exp_dir)
exp_dir

Building the MatchboxNet Model

MatchboxNet is an ASR model with a classification task - it generates one label for the entire provided audio stream. Therefore we encapsulate it inside the EncDecClassificationModel as follows.

In [ ]:

asr_model = nemo_asr.models.EncDecClassificationModel(cfg=config.model, trainer=trainer)

Training a MatchboxNet Model

As MatchboxNet is inherently a PyTorch Lightning Model, it can easily be trained in a single line - trainer.fit(model) !

Monitoring training progress

Before we begin training, let's first create a Tensorboard visualization to monitor progress

In [ ]:

try:
  from google import colab
  COLAB_ENV = True
except (ImportError, ModuleNotFoundError):
  COLAB_ENV = False

In [ ]:

# Load the TensorBoard notebook extension
if COLAB_ENV:
  %load_ext tensorboard
else:
  print("To use tensorboard, please use this notebook in a Google Colab environment.")

In [ ]:

if COLAB_ENV:
  %tensorboard --logdir {exp_dir}
else:
  print("To use tensorboard, please use this notebook in a Google Colab environment.")

Training for 5 epochs

We see below that the model begins to get modest scores on the validation set after just 5 epochs of training

In [ ]:

trainer.fit(asr_model)

Evaluation on the Test set

Lets compute the final score on the test set via trainer.test(model)

In [ ]:

trainer.test(asr_model, ckpt_path=None)

Fast Training

We can dramatically improve the time taken to train this model by using Multi GPU training along with Mixed Precision.

For multi-GPU training, take a look at the PyTorch Lightning Multi-GPU training section

For mixed-precision training, take a look at the PyTorch Lightning Mixed-Precision training section

# Mixed precision:
trainer = Trainer(amp_level='O1', precision=16)

# Trainer with a distributed backend:
trainer = Trainer(gpus=2, num_nodes=2, accelerator='ddp')

# Of course, you can combine these flags as well.

Evaluation of incorrectly predicted samples

Given that we have a trained model, which performs reasonably well, let's try to listen to the samples where the model is least confident in its predictions.

For this, we need the support of the librosa library.

NOTE: The following code depends on librosa. To install it, run the following code block first.

In [ ]:

!pip install librosa

Extract the predictions from the model

We want to possess the actual logits of the model instead of just the final evaluation score, so we can define a function to perform the forward step for us without computing the final loss. Instead, we extract the logits per batch of samples provided.

Accessing the data loaders

We can utilize the setup_test_data method in order to instantiate a data loader for the dataset we want to analyze.

For convenience, we can access these instantiated data loaders using the following accessors - asr_model._train_dl, asr_model._validation_dl and asr_model._test_dl.

In [ ]:

asr_model.setup_test_data(config.model.test_ds)
test_dl = asr_model._test_dl

Partial Test Step

Below we define a utility function to perform most of the test step. For reference, the test step is defined as follows:

    def test_step(self, batch, batch_idx, dataloader_idx=0):
        audio_signal, audio_signal_len, labels, labels_len = batch
        logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
        loss_value = self.loss(logits=logits, labels=labels)
        correct_counts, total_counts = self._accuracy(logits=logits, labels=labels)
        return {'test_loss': loss_value, 'test_correct_counts': correct_counts, 'test_total_counts': total_counts}

In [ ]:

@torch.no_grad()
def extract_logits(model, dataloader):
  logits_buffer = []
  label_buffer = []

  # Follow the above definition of the test_step
  for batch in dataloader:
    audio_signal, audio_signal_len, labels, labels_len = batch
    logits = model(input_signal=audio_signal, input_signal_length=audio_signal_len)

    logits_buffer.append(logits)
    label_buffer.append(labels)
    print(".", end='')
  print()
  
  print("Finished extracting logits !")
  logits = torch.cat(logits_buffer, 0)
  labels = torch.cat(label_buffer, 0)
  return logits, labels

In [ ]:

cpu_model = asr_model.cpu()
cpu_model.eval()
logits, labels = extract_logits(cpu_model, test_dl)
print("Logits:", logits.shape, "Labels :", labels.shape)

In [ ]:

# Compute accuracy - `_accuracy` is a PyTorch Lightning Metric !
acc = cpu_model._accuracy(logits=logits, labels=labels)
print("Accuracy : ", float(acc[0]*100))

Filtering out incorrect samples

Let us now filter out the incorrectly labeled samples from the total set of samples in the test set

In [ ]:

import librosa
import json
import IPython.display as ipd

In [ ]:

# First let's create a utility class to remap the integer class labels to actual string label
class ReverseMapLabel:
    def __init__(self, data_loader):
        self.label2id = dict(data_loader.dataset.label2id)
        self.id2label = dict(data_loader.dataset.id2label)

    def __call__(self, pred_idx, label_idx):
        return self.id2label[pred_idx], self.id2label[label_idx]

In [ ]:

# Next, let's get the indices of all the incorrectly labeled samples
sample_idx = 0
incorrect_preds = []
rev_map = ReverseMapLabel(test_dl)

# Remember, evaluated_tensor = (loss, logits, labels)
probs = torch.softmax(logits, dim=-1)
probas, preds = torch.max(probs, dim=-1)

total_count = cpu_model._accuracy.total_counts_k[0]
incorrect_ids = (preds != labels).nonzero()
for idx in incorrect_ids:
    proba = float(probas[idx][0])
    pred = int(preds[idx][0])
    label = int(labels[idx][0])
    idx = int(idx[0]) + sample_idx

    incorrect_preds.append((idx, *rev_map(pred, label), proba))

print(f"Num test samples : {total_count.item()}")
print(f"Num errors : {len(incorrect_preds)}")

# First lets sort by confidence of prediction
incorrect_preds = sorted(incorrect_preds, key=lambda x: x[-1], reverse=False)

Examine a subset of incorrect samples

Let's print out the (test id, predicted label, ground truth label, confidence) tuple of first 20 incorrectly labeled samples

In [ ]:

for incorrect_sample in incorrect_preds[:20]:
    print(str(incorrect_sample))

Define a threshold below which we designate a model's prediction as "low confidence"

In [ ]:

# Filter out how many such samples exist
low_confidence_threshold = 0.25
count_low_confidence = len(list(filter(lambda x: x[-1] <= low_confidence_threshold, incorrect_preds)))
print(f"Number of low confidence predictions : {count_low_confidence}")

Let's hear the samples which the model has least confidence in !

In [ ]:

# First let's create a helper function to parse the manifest files
def parse_manifest(manifest):
    data = []
    for line in manifest:
        line = json.loads(line)
        data.append(line)

    return data

In [ ]:

# Next, let's create a helper function to actually listen to certain samples
def listen_to_file(sample_id, pred=None, label=None, proba=None):
    # Load the audio waveform using librosa
    filepath = test_samples[sample_id]['audio_filepath']
    audio, sample_rate = librosa.load(filepath)

    if pred is not None and label is not None and proba is not None:
        print(f"Sample : {sample_id} Prediction : {pred} Label : {label} Confidence = {proba: 0.4f}")
    else:
        print(f"Sample : {sample_id}")

    return ipd.Audio(audio, rate=sample_rate)

In [ ]:

# Now let's load the test manifest into memory
test_samples = []
with open(test_dataset, 'r') as test_f:
    test_samples = test_f.readlines()

test_samples = parse_manifest(test_samples)

In [ ]:

# Finally, let's listen to all the audio samples where the model made a mistake
# Note: This list of incorrect samples may be quite large, so you may choose to subsample `incorrect_preds`
count = min(count_low_confidence, 20)  # replace this line with just `count_low_confidence` to listen to all samples with low confidence

for sample_id, pred, label, proba in incorrect_preds[:count]:
    ipd.display(listen_to_file(sample_id, pred=pred, label=label, proba=proba))

Fine-tuning on a new dataset

We currently trained our dataset on all 30/35 classes of the Google Speech Commands dataset (v1/v2).

We will now show an example of fine-tuning a trained model on a subset of the classes, as a demonstration of fine-tuning.

Preparing the data-subsets

Let's select 2 of the classes, yes and no and prepare our manifests with this dataset.

In [ ]:

import json

In [ ]:

def extract_subset_from_manifest(name: str, manifest_path: str, labels: list):
  manifest_dir = os.path.split(manifest_path)[0]
  labels = set(labels)
  manifest_values = []

  print(f"Parsing manifest: {manifest_path}")
  with open(manifest_path, 'r') as f:
    for line in f:
      val = json.loads(line)

      if val['command'] in labels:
        manifest_values.append(val)

  print(f"Number of files extracted from dataset: {len(manifest_values)}")

  outpath = os.path.join(manifest_dir, name)
  with open(outpath, 'w') as f:
    for val in manifest_values:
      json.dump(val, f)
      f.write("\n")
      f.flush()

  print("Manifest subset written to path :", outpath)
  print()

  return outpath

In [ ]:

labels = ["yes", "no"]

train_subdataset = extract_subset_from_manifest("train_subset.json", train_dataset, labels)
val_subdataset = extract_subset_from_manifest("val_subset.json", val_dataset, labels)
test_subdataset = extract_subset_from_manifest("test_subset.json", test_dataset, labels)

Saving/Restoring a checkpoint

There are multiple ways to save and load models in NeMo. Since all NeMo models are inherently Lightning Modules, we can use the standard way that PyTorch Lightning saves and restores models.

NeMo also provides a more advanced model save/restore format, which encapsulates all the parts of the model that are required to restore that model for immediate use.

In this example, we will explore both ways of saving and restoring models, but we will focus on the PyTorch Lightning method.

Saving and Restoring via PyTorch Lightning Checkpoints

When using NeMo for training, it is advisable to utilize the exp_manager framework. It is tasked with handling checkpointing and logging (Tensorboard as well as WandB optionally!), as well as dealing with multi-node and multi-GPU logging.

Since we utilized the exp_manager framework above, we have access to the directory where the checkpoints exist.

exp_manager with the default settings will save multiple checkpoints for us -

1. A few checkpoints from certain steps of training. They will have --val_loss= tags

2. A checkpoint at the last epoch of training denotes by -last.

3. If the model finishes training, it will also have a --end checkpoint.

In [ ]:

import glob

In [ ]:

print(exp_dir)

In [ ]:

# Let's list all the checkpoints we have
checkpoint_dir = os.path.join(exp_dir, 'checkpoints')
checkpoint_paths = list(glob.glob(os.path.join(checkpoint_dir, "*.ckpt")))
checkpoint_paths

In [ ]:

# We want the checkpoint saved after the final step of training
final_checkpoint = list(filter(lambda x: "-last.ckpt" in x, checkpoint_paths))[0]
print(final_checkpoint)

Restoring from a PyTorch Lightning checkpoint

To restore a model using the LightningModule.load_from_checkpoint() class method.

In [ ]:

restored_model = nemo_asr.models.EncDecClassificationModel.load_from_checkpoint(final_checkpoint)

Prepare the model for fine-tuning

Remember, the original model was trained for a 30/35 way classification task. Now we require only a subset of these models, so we need to modify the decoder head to support fewer classes.

We can do this easily with the convenient function EncDecClassificationModel.change_labels(new_label_list).

By performing this step, we discard the old decoder head, but still, preserve the encoder!

In [ ]:

restored_model.change_labels(labels)

Prepare the data loaders

The restored model, upon restoration, will not attempt to set up any data loaders.

This is so that we can manually set up any datasets we want - train and val to finetune the model, test in order to just evaluate, or all three to do both!

The entire config that we used before can still be accessed via ModelPT.cfg, so we will use it in order to set up our data loaders. This also gives us the opportunity to set any additional parameters we wish to setup!

In [ ]:

import copy

In [ ]:

train_subdataset_cfg = copy.deepcopy(restored_model.cfg.train_ds)
val_subdataset_cfg = copy.deepcopy(restored_model.cfg.validation_ds)
test_subdataset_cfg = copy.deepcopy(restored_model.cfg.test_ds)

In [ ]:

# Set the paths to the subset of the dataset
train_subdataset_cfg.manifest_filepath = train_subdataset
val_subdataset_cfg.manifest_filepath = val_subdataset
test_subdataset_cfg.manifest_filepath = test_subdataset

In [ ]:

# Setup the data loader for the restored model
restored_model.setup_training_data(train_subdataset_cfg)
restored_model.setup_multiple_validation_data(val_subdataset_cfg)
restored_model.setup_multiple_test_data(test_subdataset_cfg)

In [ ]:

# Check data loaders are correct
print("Train dataset labels :", restored_model._train_dl.dataset.labels)
print("Val dataset labels :", restored_model._validation_dl.dataset.labels)
print("Test dataset labels :", restored_model._test_dl.dataset.labels)

Setting up a new Trainer and Experiment Manager

A restored model has a utility method to attach the Trainer object to it, which is necessary in order to correctly set up the optimizer and scheduler!

Note: The restored model does not contain the trainer config with it. It is necessary to create a new Trainer object suitable for the environment where the model is being trained. The template can be replicated from any of the training scripts.

Here, since we already had the previous config object that prepared the trainer, we could have used it, but for demonstration, we will set up the trainer config manually.

In [ ]:

# Setup the new trainer object
# Let's modify some trainer configs for this demo
# Checks if we have GPU available and uses it
cuda = 1 if torch.cuda.is_available() else 0

trainer_config = OmegaConf.create(dict(
    gpus=cuda,
    max_epochs=5,
    max_steps=None,  # computed at runtime if not set
    num_nodes=1,
    accumulate_grad_batches=1,
    checkpoint_callback=False,  # Provided by exp_manager
    logger=False,  # Provided by exp_manager
    log_every_n_steps=1,  # Interval of logging.
    val_check_interval=1.0,  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
))
print(OmegaConf.to_yaml(trainer_config))

In [ ]:

trainer_finetune = pl.Trainer(**trainer_config)

Setting the trainer to the restored model

All NeMo models provide a convenience method set_trainer() in order to setup the trainer after restoration

In [ ]:

restored_model.set_trainer(trainer_finetune)

In [ ]:

exp_dir_finetune = exp_manager(trainer_finetune, config.get("exp_manager", None))

In [ ]:

exp_dir_finetune = str(exp_dir_finetune)
exp_dir_finetune

Setup optimizer + scheduler

For a fine-tuning experiment, let's set up the optimizer and scheduler!

We will use a much lower learning rate than before, and also swap out the scheduler from PolyHoldDecay to CosineDecay.

In [ ]:

optim_sched_cfg = copy.deepcopy(restored_model.cfg.optim)
# Struct mode prevents us from popping off elements from the config, so let's disable it
OmegaConf.set_struct(optim_sched_cfg, False)

In [ ]:

# Lets change the maximum learning rate to previous minimum learning rate
optim_sched_cfg.lr = 0.001

# Lets change the scheduler
optim_sched_cfg.sched.name = "CosineAnnealing"

# "power" isnt applicable to CosineAnnealing so let's remove it
optim_sched_cfg.sched.pop('power')

# "hold_ratio" isnt applicable to CosineAnnealing, so let's remove it
optim_sched_cfg.sched.pop('hold_ratio')

# Set "min_lr" to lower value
optim_sched_cfg.sched.min_lr = 1e-4

print(OmegaConf.to_yaml(optim_sched_cfg))

In [ ]:

# Now lets update the optimizer settings
restored_model.setup_optimization(optim_sched_cfg)

In [ ]:

# We can also just directly replace the config inplace if we choose to
restored_model.cfg.optim = optim_sched_cfg

Fine-tune training step

We fine-tune on the subset classification problem. Note, the model was originally trained on these classes (the subset defined here has already been trained on above).

When fine-tuning on a truly new dataset, we will not see such a dramatic improvement in performance. However, it should still converge a little faster than if it was trained from scratch.

Warning:

Output truncated. This notebook contains too many cells to display efficiently.