The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

In the root directory, the most important files are:

.
|-- common                         # Generic code for training
│   |-- fairseq                    # Parts of https://github.com/facebookresearch/fairseq
│   |-- ...
|-- inference.py                   # Evaluates trained models and measures latency
|-- scripts
│   |-- download_wav2vec2_base.sh  # Downloads pre-trained models from NGC
│   |-- finetune_base_960h.sh      # Helper script for fine-tuning with train.py
│   |-- inference.sh               # Helper script for inference.py
│   |-- pretrain_base.sh           # Helper script for pre-training with train.py
│   |-- ...
|-- train.py                       # Main pre-training and fine-tuning script
|-- utils                          # Misc standalone Python scripts
|-- wav2vec2                       # Code specific to wav2vec 2.0 model
    |-- arg_parser.py
    |-- criterion.py
    |-- logging.py
    |-- model.py
    |-- utils.py

Parameters

Parameters can be set through environment variables. The most important available parameters for scripts/pretrain_base.sh script are:

OUTPUT_DIR              directory for results, logs, and created checkpoints
                        (default: "./results/pretrain_base")
NUM_GPUS                number of GPUs to use. (default: 8)
MAX_TOKENS              upper limit for the number of tokens in a batch; changing
                        this value alters loss function consts (default: 1400000)
NUM_CONCAT_BATCHES      number of sub-batches, each with MAX_TOKENS tokens,
                        to make up one large batch (default: 8)
UPDATE_FREQ             number of grad accumulation steps before the update (default: 1)
MAX_UPDATE              training length expressed as the number of updates (default: 400000)
LEARNING_RATE           peak learning rate (default: 0.0005)
SEED                    random seed controlling model weights and data shuffling (default: disabled)
FP16                    enables mixed-precision training with float16 (default: false)
BF16                    enabled mixed-precision training with bfloat16 (default: false)
DATASET_DIR             directory with file lists (default: /datasets/LibriSpeech)
TRAIN_SUBSET            base name of the .tsv file list in the DATASET_DIR (default: "train-full-960")
VALID_SUBSET            base name of the validation .tsv file list in the DATASET_DIR (default: "dev-other")
SAVE_FREQUENCY          frequency of saving checkpoints to disk (default: 1)
HOURGLASS_CONFIG        configuration of Hourglass Transformer; refer to the section
                        below for details (default: "[2,(8,4),2]")

In addition, important parameters for scripts/finetune_base_960h.sh script are:

PRETRAINED_MODEL        a path to a pre-trained model checkpoint for fine-tuning
                        (default: "./results/pretrain_base/wav2vec2_update400000.pt")
FREEZE_FINETUNE_UPDATES freeze wav2vec 2.0 encoder for an initial number of steps and train only
                        the output linear projection (default: 0)

Below we present more details on how to set crucial parameters.

Adjusting batch size and the number of GPUs

Every training recipe assumes a constant world size, and variables need to be adjusted to maintain that world size, for example, NUM_GPUS x NUM_CONCAT_BATCHES x UPDATE_FREQUENCY = 64 for pre-training of the Base model:

• first, set NUM_GPUS to the number of available GPUs,
• then, adjust NUM_CONCAT_BATCHES to a high value that does not cause out-of-memory errors
• finally, adjust the update frequency that controls gradient accumulation, to maintain the effective world size.

NUM_CONCAT_BATCHES controls the number of sub-batches that are forwarded through the model, each with --max_tokens tokens. In the case of out-of-memory errors, it has to be lowered. With Hourglass Transformer and mixed-precision training, the model should fit within 12GB of GPU memory on the lowest NUM_CONCAT_BATCHES=1 setting.

Adjusting mixed precision

By default, the model is trained in TF32 (A100 GPUs) or FP32 (V100 and older GPUs). Mixed-precision training can be performed in float16 or bfloat16 precisions. Training in bfloat16 is more stable and requires less stabilizing casts to FP32; thus, it is a bit faster. It is supported on the hardware level in NVIDIA Ampere and newer architectures. Scripts scripts/pretrain_base.sh and scripts/finetune_base_960h.sh provide env vars for setting appropriate casting flags. In order to benefit from mixed-precision training, set either BF16=true or FP16=true, depending on the architecture of the GPU.

Adjusting Hourglass Transformer

The Hourglass Transformer architecture is configurable by four parameters:

• the number of initial transformer layers,
• the number of middle transformer layers that process the downsampled signal,
• downsampling rate,
• the number of output transformer layers.

These are expressed in that exact order by a Python list without whitespace. For instance, the default setting is HOURGLASS_CONFIG="[2,(8,4),2]". It uses 12 layers in total (two initial, eight middle with a downsampling rate 4, and two output layers).

During fine-tuning, the same architecture as during pre-training has to be set.

Command-line options

To view the full list of available options and their descriptions, use the -h or --help command-line option, for example: python train.py -h. Most of the command-line options are a subset of those from the original Fairseq wav2vec 2.0 codebase.

Getting the data

The wav2vec 2.0 model described in the paper was pre-trained on either the LibriSpeech or LibriVox datasets. We publish recipes for training on pre-training and fine-tuning on the LibriSpeech dataset. The dev-other subset is used as a validation dataset, and test-other is used as a testing dataset.

The ./scripts/download_ls_dataset.sh [TARGET_DATA_DIR] script downloads and extracts the LibriSpeech dataset to the directory of choice, by default /datasets/LibriSpeech if the argument is omitted.

The ./scripts/generate_ls_filelists.sh [SOURCE_DATA_DIR] [TARGET_FILELISTS_DIR] script prepares filelists and collect transcriptions. Again, positional arguments are optional and default to /datasets/LibriSpeech.

Dataset guidelines

LibriSpeech data is kept at the default sampling rate of 16 kHz. The model works with either .wav or .flac files. Both are lossless, with .flac being more efficient in terms of storage but requiring extra computation during training. Files are listed in .tsv filelists. The first row is the top-level directory, and subsequent lines listths to files and a number of samples delimited by tab:

/datasets/LibriSpeech/test-other
367/293981/367-293981-0017.flac\t46560
367/293981/367-293981-0009.flac\t52720
...

The .ltr files, generated alongside .tsv filelists, hold character-level transcriptions for filelists with the same basename. Filelists and transcription lists should list samples in matching order.

A N D | A | V E R Y | R E S P E C T A B L E | O N E | S A I D | T H E | I N N K E E P E R |
T H E | O F F I C E R | T U R N E D | T O | H I M | A N D | S A I D | W E L L | H O W | G O E S | I T | G O O D | M A N |
...

Finally, generate a dict.ltr.txt dictionary using training .ltr transcripts:

python utils/generate_dictionary.py /my/dataset/path/train.ltr /my/dataset/path/dict.ltr.txt

Multi-dataset

In order to train on multiple datasets, prepare a filelist and transcription list with all files from those datasets. Refer to scripts/generate_filelists.sh for an example of concatenating LibriSpeech training filelists.

Training process

Training of wav2vec 2.0 is performed in two stages: unsupervised pre-training and supervised fine-tuning. Both are performed with the train.py script.

Pre-training The scripts/pretrain_base.sh script sets command-line arguments for train.py and runs a job on a single node that trains the wav2vec 2.0 model from scratch. Key variables can be conveniently changed via env variables.

Fine-tuning The scripts/finetune_base_960h.sh script sets command-line arguments for train.py and runs a job on a single node that fine-tunes a pre-trained wav2vec 2.0 model. Key variables can be conveniently changed via env variables. Note that a checkpoint trained with Fairseq can be loaded and fine-tuned using this repository.

Apart from the arguments as listed in the Parameters section, by default both training scripts:

• Run on eight GPUs with at least 80GB of memory with increased batch size, so that gradient accumulation is not necessary
• Use TF32 precision (A100 GPU) or FP32 (other GPUs)
• Use Hourglass Transformer architecture with shortening factor of 4
• Train on 960 hours of LibriSpeech training data and evaluate on the dev-other subset
• Remove old checkpoints and preserve milestone checkpoints automatically
• Maintain a separate checkpoint with the lowest WER on the dev set
• Create a DLLogger log file and a TensorBoard log
• Set the remaining parameters according to the recipes published with the original paper

The current training setup recreates WER Results published in the original paper, while significantly lowering the time and memory required for training.

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 fine-tuned wav2vec 2.0 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 test-other dataset and prints out the final word error rate
• Uses a batch size of 8
• Creates a log file with progress and results, which will be stored in the results folder
• Does greedy decoding and optionally saves the transcriptions in the results folder
• Has the option to save the model output tensors for more complex decoding, for example, beam search

To view all available options for inference, run python inference.py --help