Exportable mel spectrogram preprocessor
State: closed | Pull request: yes
Labels: name: ASR

name: NLP

# What does this PR do ?

`AudioToMelSpectrogramPreprocessor` accepts a bool argument `use_torchaudio` which, if True, switches the `featurizer` to a `torchaudio`-based extractor which produces the same features but with an exportable graph.

New preprocessor mimics the old implementation sufficiently well to swap out the preprocessor of pre-trained models and export them.

Preprocessor can be exported to JIT; ONNX is blocked by https://github.com/pytorch/pytorch/issues/81075 

**Collection**: ASR

# Changelog 
Add an option to `AudioToMelSpectrogramPreprocessor` which can alter `featurizer`. 

Add a class `FilterbankFeaturesTA` which is analogous to `FilterbankFeatures`

# Usage

The following script will
1. Load a pre-trained English Conformer
2. Create a copy of the model's preprocessor, but with a `torchaudio` backend
3. Compare new and old preprocessor outputs
4. Export the preprocessor to JIT
5. Compare JIT and PyTorch outputs
6. Swap the pre-trained model's preprocessor to the new one, check WER matches old

```python

from pathlib import Path

import torch
import hydra
from pytorch_lightning import seed_everything
from omegaconf import open_dict

from nemo.utils import logging
from nemo.utils.nemo_logging import Logger
from nemo.collections.asr.metrics.wer import word_error_rate
from nemo.collections.asr.models import ASRModel
from nemo.collections.asr.modules.audio_preprocessing import AudioToMelSpectrogramPreprocessor


# Optionally use a seed
# seed_everything(42)

logging.set_verbosity(Logger.CRITICAL)

# Get a pre-trained ASR model to compare preprocessors. Any model with a mel spec extractor should work.
m: ASRModel = ASRModel.from_pretrained("stt_en_conformer_ctc_small", map_location=torch.device("cpu"))
m.eval()
old_preprocessor = m.preprocessor

# Extract preprocessor config and set the flag to use the torchaudio-based extractor; keep all other arguments the same
new_config = m.cfg.preprocessor
with open_dict(new_config):
    new_config.use_torchaudio = True
# Instantiate an instance that uses torchaudio on the backend
new_preprocessor: AudioToMelSpectrogramPreprocessor = hydra.utils.instantiate(config=new_config)
new_preprocessor.eval()
print(f"New preprocessor featurizer type: {type(new_preprocessor.featurizer)}")

# Export the torchaudio preprocessor and load it back in as a `ScriptModule`.
new_preprocessor.export("tmp.pt")
jit_preprocessor = torch.jit.load("tmp.pt")

# Generate random input
batch_size = 4
max_length = 16000
signals = torch.randn(size=[batch_size, max_length])
lengths = torch.randint(low=200, high=max_length, size=[batch_size])
lengths[0] = max_length

# Extract features with all preprocessors
old_feats, old_feat_lens = old_preprocessor(input_signal=signals, length=lengths)
new_feats, new_feat_lens = new_preprocessor(input_signal=signals, length=lengths)
jit_feats, jit_feat_lens = jit_preprocessor(input_signal=signals, length=lengths)

# Make sure new output matches old output
# Need to relax the tolerance from defaults. We will check WER also, as an alternative verification of correctness.
rel_tolerance = 1e-2
abs_tolerance = 1e-4
torch.testing.assert_allclose(actual=new_feats, expected=old_feats, atol=abs_tolerance, rtol=rel_tolerance)
# Zero tolerance for integer lengths.
torch.testing.assert_allclose(actual=new_feat_lens, expected=old_feat_lens, atol=0, rtol=0)

print(f"Output comparison passed with relative tolerance {rel_tolerance} and absolute tolerance {abs_tolerance}.")

# Make sure JIT output matches PyTorch output
# Keep tolerance at defaults for JIT comparison
torch.testing.assert_allclose(actual=jit_feats, expected=new_feats)
torch.testing.assert_allclose(actual=jit_feat_lens, expected=new_feat_lens, atol=0, rtol=0)
print(f"Jit comparison passed with default tolerance.")

# To run a WER check you'll need to comment out some assumptions in the CTC transcribe method, as addressed in https://github.com/NVIDIA/NeMo/pull/2762 
# print("Testing WER with old/new preprocessor with some LibriSpeech data")
# We only need audio files; we're comparing model outputs to each other, not references
# dev_other_dir = "/path/to/LibriSpeech/dev-other"
# num_files_to_use = 100
# audio_files = [str(x) for x in Path(dev_other_dir).rglob("*.flac")]
# audio_files = audio_files[:num_files_to_use]
# print("Transcribing with the baseline model")
# old_output = m.transcribe(audio_files)
# m.preprocessor = new_preprocessor
# print("Transcribing after switching the preprocessor")
# new_output = m.transcribe(audio_files)
# wer = word_error_rate(hypotheses=new_output, references=old_output)
# print(f"WER with {len(audio_files)} audio files using old vs. new preprocessor is {wer}")
```

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] Did you write any new necessary tests?
- [x] Did you add or update any necessary documentation?
- [x] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to: this issue comes up once in a while and is generally brushed aside.