The repository contains a PyTorch reproduction of the TM-CTC model from the Deep Audio-Visual Speech Recognition paper. We train three models - Audio-Only (AO), Video-Only (VO) and Audio-Visual (AV), on the LRS2 dataset for the speech-to-text transcription task.
System packages:
ffmpeg==2.8.15
python==3.6.9
Python packages:
editdistance==0.5.3
matplotlib==3.1.1
numpy==1.18.1
opencv-python==4.2.0
pytorch==1.2.0
scipy==1.3.1
tqdm==4.42.1
CUDA 10.0 (if NVIDIA GPU is to be used):
cudatoolkit==10.0
The structure of the audio_only
, video_only
and audio_visual
directories is as follows:
/checkpoints
: Temporary directory to store intermediate model weights and plots while training. Gets automatically created.
/data
: Directory containing the LRS2 Main and Pretrain dataset class definitions and other required data-related utility functions.
/final
: Directory to store the final trained model weights and plots. If available, place the pre-trained model weights in the models
subdirectory.
/models
: Directory containing the class definitions for the models.
/utils
: Directory containing function definitions for calculating CER/WER, greedy search/beam search decoders and preprocessing of data samples. Also contains functions to train and evaluate the model.
checker.py
: File containing checker/debug functions for testing all the modules and the functions in the project as well as any other checks to be performed.
config.py
: File to set the configuration options and hyperparameter values.
demo.py
: Python script for generating predictions with the specified trained model for all the data samples in the specified demo directory.
preprocess.py
: Python script for preprocessing all the data samples in the dataset.
pretrain.py
: Python script for pretraining the model on the pretrain set of the LRS2 dataset using curriculum learning.
test.py
: Python script to test the trained model on the test set of the LRS2 dataset.
train.py
: Python script to train the model on the train set of the LRS2 dataset.
We provide Word Error Rate (WER) achieved by the models on the test set of the LRS2 dataset with both Greedy Search and Beam Search (with Language Model) decoding techniques. We have tested in cases of clean audio and noisy audio (0 dB SNR). We also give WER in cases where only one of the modalities is used in the Audio-Visual model.
Operation Mode | AO/VO Model | AV Model | ||
---|---|---|---|---|
Greedy | Beam (+LM) |
Greedy | Beam (+LM) | |
Clean Audio | ||||
AO | 11.4% | 8.3% | 12.0% | 8.2% |
VO | 61.8% | 55.3% | 56.3% | 49.2% |
AV | - | - | 10.3% | 6.8% |
Noisy Audio | ||||
AO | 62.5% | 54.0% | 59.0% | 50.7% |
AV | - | - | 29.1% | 22.1% |
Download the pre-trained weights for the Visual Frontend, AO, VO, AV and Language model from here.
Once the Visual Frontend and Language Model weights are downloaded, place them in a folder and add their paths in the config.py
file. Place the weights of AO, VO and AV model in their corresponding /final/models
directory.
If planning to train the models, download the complete LRS2 dataset from here or in cases of custom datasets, have the specifications and folder structure similar to LRS2 dataset.
Steps have been provided to either train the models or to use the trained models directly for inference:
Set the configuration options in the config.py
file before each of the following steps as required. Comments have been provided for each option. Also, check the Training Details section below as a guide for training the models from scratch.
-
Run the
preprocess.py
script to preprocess and generate the required files for each sample. -
Run the
pretrain.py
script for one iteration of curriculum learning. Run it multiple times, each time changing thePRETRAIN_NUM_WORDS
argument in theconfig.py
file to perform multiple iterations of curriculum learning. -
Run the
train.py
script to finally train the model on the train set. -
Once the model is trained, run the
test.py
script to obtain the performance of the trained model on the test set. -
Run the
demo.py
script to use the model to make predictions for each sample in a demo directory. Read the specifications for the sample in thedemo.py
file.
-
Set the configuration options in the
config.py
file. Comments have been provided for each option. -
Run the
demo.py
script to use the model to make predictions for each sample in a demo directory. Read the specifications for the sample in thedemo.py
file.
-
We perform iterations of Curriculum Learning by changing the
PRETRAIN_NUM_WORDS
config option. The number of words used in each iteration of curriculum learning is as follows: 1,2,3,5,7,9,13,17,21,29,37, i.e., 11 iterations in total. -
During Curriculum Learning, the minibatch size (default=32) is reduced by half each time we hit an Out Of Memory error.
-
In each iteration, the training is terminated forcefully once the validation set WER flattens. We also make sure that the Learning Rate has decreased to the minimum value before terminating the training.
-
We train the AO and VO models first. We then initialize the AV model with weights from the trained AO and VO models as follows: AO Audio Encoder → AV Audio Encoder, VO Video Encoder → AV Video Encoder, VO Video Decoder → AV Joint Decoder.
-
The weights of the Audio and Video Encoders are fixed during AV model pretraining. The complete AV model is trained on the train set after the pretraining is complete.
-
We have used a GPU with 11 GB memory for our training. Each model took around 7 days for complete training.
-
The pre-trained weights of the Visual Frontend and the Language Model have been obtained from Afouras T. and Chung J, Deep Lip Reading: a comparison of models and an online application, 2018 GitHub repository.
-
The CTC beam search implementation is adapted from Harald Scheidl, CTC Decoding Algorithms GitHub repository.