Deep Audio-Visual Speech Recognition
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.
Requirements
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
Project Structure
The structure of the audio_only, video_only and audio_visual directories is as follows:
Directories
/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.
Files
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.
Results
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% |
Pre-trained Weights
Download the pre-trained weights for the Visual Frontend from here and for the 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.
For the pre-trained weights of the AO, VO and AV models, please send an email at smeet.shah.c2020<AT>iitbombay.org from your institutional email ID. Place the weights of each model in the corresponding /final/models directory.
Note:
- Replace
<AT>with@in the email ID while sending an email. It may take few days to reply back. Please have patience. - Do NOT open issues on GitHub requesting pre-trained weights. Such issues may be deleted. I am sharing weights only via replies to email requests from valid institutional email IDs.
- The link for these weights has also been updated. If required, please resend an email for new link.
How To Use
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:
Training
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.pyscript to preprocess and generate the required files for each sample. -
Run the
pretrain.pyscript for one iteration of curriculum learning. Run it multiple times, each time changing thePRETRAIN_NUM_WORDSargument in theconfig.pyfile to perform multiple iterations of curriculum learning. -
Run the
train.pyscript to finally train the model on the train set. -
Once the model is trained, run the
test.pyscript to obtain the performance of the trained model on the test set. -
Run the
demo.pyscript to use the model to make predictions for each sample in a demo directory. Read the specifications for the sample in thedemo.pyfile.
Inference
-
Set the configuration options in the
config.pyfile. Comments have been provided for each option. -
Run the
demo.pyscript to use the model to make predictions for each sample in a demo directory. Read the specifications for the sample in thedemo.pyfile.
Important Training Details
-
We perform iterations of Curriculum Learning by changing the
PRETRAIN_NUM_WORDSconfig 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.
References
-
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.
PS: Please do not hesitate to raise an issue in case of any bugs/doubts/suggestions. However, please be patient if it takes some time to reply. Happy Open Source-ing !!