The Jasper model is an end-to-end neural acoustic model for automatic speech recognition (ASR).
The following sections provide greater details of the dataset, running training and inference, and getting training and inference results.
Scripts and sample code
In the root directory, the most important files are:
jasper
|-- common # data pre-processing, logging, etc.
|-- configs # model configurations
|-- Dockerfile # container with the basic set of dependencies to run Jasper
|-- inference.py # entry point for inference
|-- jasper # model-specific code
|-- notebooks # jupyter notebooks and example audio files
|-- scripts # one-click scripts required for running various supported functionalities
│ |-- docker # contains the scripts for building and launching the container
│ |-- download_librispeech.sh # downloads LibriSpeech dataset
│ |-- evaluation.sh # runs evaluation using the `inference.py` script
│ |-- inference_benchmark.sh # runs the inference benchmark using the `inference_benchmark.py` script
│ |-- inference.sh # runs inference using the `inference.py` script
│ |-- preprocess_librispeech.sh # preprocess LibriSpeech raw data files for training and inference
│ |-- train_benchmark.sh # runs the training performance benchmark using the `train.py` script
│ |-- train.sh # runs training using the `train.py` script
|-- train.py # entry point for training
|-- triton # example of inference using Triton Inference Server
|-- utils # data downloading and common routines
Parameters
Parameters could be set as env variables, or passed as positional arguments.
The complete list of available parameters for scripts/train.sh script contains:
DATA_DIR: directory of dataset. (default: '/datasets/LibriSpeech')
MODEL_CONFIG: relative path to model configuration. (default: 'configs/jasper10x5dr_speedp-online_speca.yaml')
OUTPUT_DIR: directory for results, logs, and created checkpoints. (default: '/results')
CHECKPOINT: a specific model checkpoint to continue training from. To resume training from the last checkpoint, see the RESUME option.
RESUME: resume training from the last checkpoint found in OUTPUT_DIR, or from scratch if there are no checkpoints (default: true)
CUDNN_BENCHMARK: boolean that indicates whether to enable cudnn benchmark mode for using more optimized kernels. (default: true)
NUM_GPUS: number of GPUs to use. (default: 8)
AMP: if set to `true`, enables automatic mixed precision (default: false)
BATCH_SIZE: effective data batch size. The real batch size per GPU might be lower, if gradient accumulation is enabled (default: 64)
GRAD_ACCUMULATION_STEPS: number of gradient accumulation steps until optimizer updates weights. (default: 2)
LEARNING_RATE: initial learning rate. (default: 0.01)
MIN_LEARNING_RATE: minimum learning rate, despite LR scheduling (default: 1e-5)
LR_POLICY: how to decay LR (default: exponential)
LR_EXP_GAMMA: decay factor for the exponential LR schedule (default: 0.981)
EMA: decay factor for exponential averages of checkpoints (default: 0.999)
SEED: seed for random number generator and used for ensuring reproducibility. (default: 0)
EPOCHS: number of training epochs. (default: 440)
WARMUP_EPOCHS: number of initial epoch of linearly increasing LR. (default: 2)
HOLD_EPOCHS: number of epochs to hold maximum LR after warmup. (default: 140)
SAVE_FREQUENCY: number of epochs between saving the model to disk. (default: 10)
EPOCHS_THIS_JOB: run training for this number of epochs. Does not affect LR schedule like the EPOCHS parameter. (default: 0)
DALI_DEVICE: device to run the DALI pipeline on for calculation of filterbanks. Valid choices: cpu, gpu, none. (default: gpu)
PAD_TO_MAX_DURATION: pad all sequences with zeros to maximum length. (default: false)
EVAL_FREQUENCY: number of steps between evaluations on the validation set. (default: 544)
PREDICTION_FREQUENCY: the number of steps between writing a sample prediction to stdout. (default: 544)
TRAIN_MANIFESTS: lists of .json training set files
VAL_MANIFESTS: lists of .json validation set files
The complete list of available parameters for scripts/inference.sh script contains:
DATA_DIR: directory of dataset. (default: '/datasets/LibriSpeech')
MODEL_CONFIG: model configuration. (default: 'configs/jasper10x5dr_speedp-online_speca.yaml')
OUTPUT_DIR: directory for results and logs. (default: '/results')
CHECKPOINT: model checkpoint path. (required)
DATASET: name of the LibriSpeech subset to use. (default: 'dev-clean')
LOG_FILE: path to the DLLogger .json logfile. (default: '')
CUDNN_BENCHMARK: enable cudnn benchmark mode for using more optimized kernels. (default: false)
MAX_DURATION: filter out recordings shorter then MAX_DURATION seconds. (default: "")
PAD_TO_MAX_DURATION: pad all sequences with zeros to maximum length. (default: false)
PAD_LEADING: pad every batch with leading zeros to counteract conv shifts of the field of view. (default: 16)
NUM_GPUS: number of GPUs to use. Note that with > 1 GPUs WER results might be inaccurate due to the batching policy. (default: 1)
NUM_STEPS: number of batches to evaluate, loop the dataset if necessary. (default: 0)
NUM_WARMUP_STEPS: number of initial steps before measuring performance. (default: 0)
AMP: enable FP16 inference with AMP. (default: false)
BATCH_SIZE: data batch size. (default: 64)
EMA: Attempt to load exponentially averaged weights from a checkpoint. (default: true)
SEED: seed for random number generator and used for ensuring reproducibility. (default: 0)
DALI_DEVICE: device to run the DALI pipeline on for calculation of filterbanks. Valid choices: cpu, gpu, none. (default: gpu)
CPU: run inference on CPU. (default: false)
LOGITS_FILE: dump logit matrices to a file. (default: "")
PREDICTION_FILE: save predictions to a file. (default: "${OUTPUT_DIR}/${DATASET}.predictions")
The complete list of available parameters for scripts/evaluation.sh is the same as for scripts/inference.sh except for the few default changes.
PREDICTION_FILE: (default: "")
DATASET: (default: "test-other")
The scripts/inference_benchmark.sh script pads all input to a fixed duration and computes the mean, 90%, 95%, 99% percentile of latency for the specified number of inference steps. Latency is measured in milliseconds per batch. The scripts/inference_benchmark.sh measures latency for a single GPU and loops over a number of batch sizes and durations. It extends scripts/inference.sh, and changes the defaults with:
BATCH_SIZE_SEQ: batch sizes to measure on. (default: "1 2 4 8 16")
MAX_DURATION_SEQ: input durations (in seconds) to measure on (default: "2 7 16.7")
CUDNN_BENCHMARK: (default: true)
PAD_TO_MAX_DURATION: (default: true)
PAD_LEADING: (default: 0)
NUM_WARMUP_STEPS: (default: 10)
NUM_STEPS: (default: 500)
DALI_DEVICE: (default: cpu)
The scripts/train_benchmark.sh script pads all input to the same length according to the input argument MAX_DURATION and measures average training latency and throughput performance. Latency is measured in seconds per batch, throughput in sequences per second.
Training performance is measured with on-line speed perturbation and cuDNN benchmark mode enabled.
The script scripts/train_benchmark.sh loops over a number of batch sizes and GPU counts.
It extends scripts/train.sh, and the complete list of available parameters for scripts/train_benchmark.sh script contains:
BATCH_SIZE_SEQ: batch sizes to measure on. (default: "1 2 4 8 16")
NUM_GPUS_SEQ: number of GPUs to run the training on. (default: "1 4 8")
MODEL_CONFIG: (default: "configs/jasper10x5dr_speedp-online_train-benchmark.yaml")
TRAIN_MANIFESTS: (default: "$DATA_DIR/librispeech-train-clean-100-wav.json")
RESUME: (default: false)
EPOCHS_THIS_JOB: (default: 2)
EPOCHS: (default: 100000)
SAVE_FREQUENCY: (default: 100000)
EVAL_FREQUENCY: (default: 100000)
GRAD_ACCUMULATION_STEPS: (default: 1)
PAD_TO_MAX_DURATION: (default: true)
EMA: (default: 0)
Command-line options
To see the full list of available options and their descriptions, use the -h or --help command-line option with the Python file, for example:
python train.py --help
python inference.py --help
Getting the data
The Jasper model was trained on the LibriSpeech dataset. We use the concatenation of train-clean-100, train-clean-360 and train-other-500 for training and dev-clean for validation.
This repository contains the scripts/download_librispeech.sh and scripts/preprocess_librispeech.sh scripts which will automatically download and preprocess the training, test and development datasets. By default, data will be downloaded to the /datasets/LibriSpeech directory, a minimum of 250GB free space is required for download and preprocessing, the final preprocessed dataset is approximately 100GB. With offline speed perturbation, the dataset will be about 3x larger.
Dataset guidelines
The scripts/preprocess_librispeech.sh script converts the input audio files to WAV format with a sample rate of 16kHz, target transcripts are stripped from whitespace characters, then lower-cased. For train-clean-100, train-clean-360 and train-other-500. It can optionally create speed perturbed versions with rates of 0.9 and 1.1 for data augmentation. In the current version, those augmentations are applied on-line with the DALI pipeline without any impact on training time.
After preprocessing, the script creates JSON files with output file paths, sample rate, target transcript and other metadata. These JSON files are used by the training script to identify training and validation datasets.
The Jasper model was tuned on audio signals with a sample rate of 16kHz, if you wish to use a different sampling rate then some hyperparameters might need to be changed - specifically window size and step size.
Training process
The training is performed using train.py script along with parameters defined in scripts/train.sh
The scripts/train.sh script runs a job on a single node that trains the Jasper model from scratch using LibriSpeech as training data. To make training more efficient, we discard audio samples longer than 16.7 seconds from the training dataset, the total number of these samples is less than 1%. Such filtering does not degrade accuracy, but it allows us to decrease the number of time steps in a batch, which requires less GPU memory and increases training speed.
Apart from the default arguments as listed in the Parameters section, by default the training script:
- Runs on 8 GPUs with at least 32GB of memory and training/evaluation batch size 64, split over two gradient accumulation steps
- Uses TF32 precision (A100 GPU) or FP32 (other GPUs)
- Trains on the concatenation of all 3 LibriSpeech training datasets and evaluates on the LibriSpeech dev-clean dataset
- Maintains an exponential moving average of parameters for evaluation
- Has cudnn benchmark enabled
- Runs for 440 epochs
- Uses an initial learning rate of 0.01 and an exponential learning rate decay
- Saves a checkpoint every 10 epochs
- Automatically removes old checkpoints and preserves milestone checkpoints
- Runs evaluation on the development dataset every 544 iterations and at the end of training
- Maintains a separate checkpoint with the lowest WER on development set
- Prints out training progress every iteration to stdout
- Creates a DLLogger logfile and a Tensorboard log
- Calculates speed perturbation on-line during training
- Uses SpecAugment in data pre-processing
- Filters out audio samples longer than 16.7 seconds
- Pads each batch so its length would be divisible by 16
- Uses masked convolutions and dense residuals as described in the paper
- Uses weight decay of 0.001
- Uses Novograd as optimizer with betas=(0.95, 0)
Enabling AMP permits batch size 64 with one gradient accumulation step. In the current setup it will improve upon the greedy WER Results of the Jasper paper on a DGX-1 with 32GB V100 GPUs.
Inference process
Inference is performed using the inference.py script along with parameters defined in scripts/inference.sh.
The scripts/inference.sh script runs the job on a single GPU, taking a pre-trained Jasper model checkpoint and running it on the specified dataset.
Apart from the default arguments as listed in the Parameters section by default the inference script:
- Evaluates on the LibriSpeech dev-clean dataset
- Uses a batch size of 64
- Runs for 1 epoch and prints out the final word error rate
- Creates a log file with progress and results which will be stored in the results folder
- Pads each batch so its length would be divisible by 16
- Does not use data augmentation
- Does greedy decoding and saves the transcription in the results folder
- Has the option to save the model output tensors for more complex decoding, for example, beam search
- Has cudnn benchmark disabled
Evaluation process
Evaluation is performed using the inference.py script along with parameters defined in scripts/evaluation.sh.
The setup is similar to scripts/inference.sh, with two differences:
- Evaluates the LibriSpeech test-other dataset
- Model outputs are not saved
Deploying Jasper using Triton Inference Server
The NVIDIA Triton Inference Server provides a datacenter and cloud inferencing solution optimized for NVIDIA GPUs. The server provides an inference service via an HTTP or gRPC endpoint, allowing remote clients to request inferencing for any number of GPU or CPU models being managed by the server. More information on how to perform inference using Triton Inference Server with different model backends can be found in the subfolder ./https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechRecognition/Jasper/triton/README.md