NVIDIA Shows How To Build AI Models At Scale With PyTorch Lightning
Deploying predictive AI models across a business is no easy feat. The effectiveness of AI depends on the quality of the underlying models. Therefore, it becomes crucial for data scientists and academic researchers to quickly build models with various parameters and identify the most effective ones to deploy them easily and scale them seamlessly.
In its latest blog post, NVIDIA researchers showed how to build speech models with PyTorch Lighting on CPU-powered AWS instances (Grid). PyTorch Lightning is a lightweight PyTorch wrapper designed to make high performance AI research simple. It is an organised PyTorch, which allows users to train their models on CPU, GPUs, or multiple nodes without changing any code.
Grid, which runs on AWS, supports Lightning and classic machine learning frameworks such as TensorFlow, Keras, PyTorch, Sci-Kit, and others. It also helps users to scale the training of models from the NGC catalogue. NGC catalogue is a curated set of GPU-optimised containers for deep learning, visualisation, and high-performance computing (HPC).
PyTorch lightning software and developer environment is available on NGC Catalog. Also, check out GitHub to get started with Grid, NGC, PyTorch Lightning here.
Training AI Models
For building speech models, NVIDIA researchers have used ASR, which transcribes spoken language to text. ASR is a critical component of speech-to-text systems. So, when training ASR models, the goal is to generate text from a given audio input that reduces the word error rate (WER) metric on human transcribed speech. The NGC catalogue contains SOTA pretrained models of ASR.
Further, they use Grid sessions, NVIDIA NeMo, and PyTorch Lightning to fine-tune these models on the AN4 dataset, aka Alphanumeric dataset. Collected and published by Carnegie Mellon University, the AN4 dataset consists of recordings of people spelling out addresses, names, phone numbers, etc.
Here are the key steps to follow when building speech models:
Create a Grid session optimised for Lightning and pretrained NGC models
Clone the ASR demo repo and open the tutorial notebook
Install NeMo ASR dependencies
Convert and visualise the AN4 dataset (Spectrograms and Mel spectrograms)
Users can run Grid sessions on the same hardware they need to scale while providing them with pre-configured environments to iterate the ML process faster. Here, sessions are linked to GitHub, loaded with ‘JupyterHub,’ and can be accessed through SSH and IDE without installing.
Check out the Grid Session tour here. (requires Grid.ai account)
Clone ASR demo repo & open tutorial notebook
Once you have a developer environment optimised for PyTorch Lightning, the next step is to clone the NGC-Lightning-Grid-Workshop repo. After this, the user can open up the notebook to fine-tune the NGC hosted model with NeMo and PyTorch Lightning.
Install NeMo ASR dependencies
Install all the session dependencies by running tools like PyTorch Lightning and NeMo, and process the AN4 dataset. Then, run the first cell in the tutorial notebook, which runs the following bash commands to install the dependencies.
Convert and Visualise the AN4 Dataset
The AN4 dataset contains raw Sof audio files. Convert them to the Wav format so that you can use NeMo audio processing.
Once processed, you can then visualise the audio example as images of the audio waveform. The below image shows the activity in the waveform that corresponds to each letter in the audio. Each spoken letter has a different “shape.” Interestingly, the last two blobs look relatively similar because they are both the letter N.
The audio waveform of the sample example (Source: NVIDIA)