wav2vec 2.0 for PyTorch

NVIDIA Deep Learning Examples

Resource

NVIDIA Deep Learning Examples

wav2vec 2.0 for PyTorch

A framework for self-supervised learning of speech representations which masks latent representations of the raw waveform and solves a contrastive task over quantized speech representations.

To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the wav2vec 2.0 model on the LibriSpeech dataset. For the specifics concerning training and inference, refer to the Advanced section.

Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/SpeechRecognition/wav2vec2

Build the 22.11-py3 PyTorch NGC container and start an interactive session to run training/inference. DATASET_DIR on the host will be mounted as /datasets inside the container.
```
bash scripts/docker/build.sh
DATASET_DIR=[PATH] bash scripts/docker/run.sh
```
Download and preprocess the dataset. The dataset size is about 70GB and this step could take up to a few hours to complete.
```
bash scripts/download_data.sh
```
Generate filelists.
```
bash scripts/generate_filelists.sh
```

Start pre-training.

NUM_GPUS=[NUM] UPDATE_FREQUENCY=[NUM] NUM_CONCAT_BATCHES=[NUM] BF16=[true|false] FP16=[true|false] \
    bash scripts/pretrain_base.sh

Adjust the variables to maintain NUM_GPUS x NUM_CONCAT_BATCHES x UPDATE_FREQUENCY = 64. For more details, refer to Adjusting batch size and the number of GPUs and Adjusting mixed precision.

For instance:

# Mixed precision training on 4x A100 40GB
NUM_GPUS=4 NUM_CONCAT_BATCHES=8 UPDATE_FREQUENCY=2 BF16=true bash scripts/pretrain_base.sh

Start fine-tuning.

PRETRAINED_MODEL=[PATH] NUM_GPUS=[NUM] UPDATE_FREQUENCY=[NUM] BF16=[true|false] FP16=[true|false] \
    bash scripts/finetune_base_960h.sh

Adjust the variables to maintain NUM_GPUS x NUM_CONCAT_BATCHES x UPDATE_FREQUENCY = 8.

Start inference/predictions.

FINETUNED_MODEL=[PATH] BF16=[true|false] FP16=[true|false] BATCH_SIZE=[NUM] bash scripts/inference.sh

Now that you have your model trained and evaluated, you can choose to compare your training results with our Training accuracy results. You can also choose to benchmark your performance to Training performance benchmark or Inference performance benchmark. Following the steps in these sections ensures you achieve the same accuracy and performance results as stated in the Results section.