Wav2vec2.0 model support in torchaudio

mthrok · May 13, 2021, 8:47pm

Hi PyTorch dev community

I posted the project about supporting wav2vec2.0 in torchaudio.

github.com/pytorch/audio

wav2vec2.0 model support in torchaudio

opened 08:36PM - 13 May 21 UTC

mthrok

RFC

# wav2vec2.0 model with TorchScript support Following the popularity of wav2v…ec2.0 models, we plan to add wav2vec2.0 model in torchaudio. ## What are we going to add? 1. Model definitions in `torchaudio`, that supports TorchScript, specifically `torch.jit.script`. 2. Conversion of pretrained models from Hugging Face's `transformers` (and `fairseq`).* \* For importing models of these libraries, you will need to install them, but they will be optional dependencies of `torchaudio`. ## What is the value of this addition? With the support for TorchScript, one can deploy wav2vec2.0 models into non-Python environments, such as C++ applications and mobiles. There is [an example of running wav2vec2.0 in mobile](https://github.com/pytorch/android-demo-app/tree/master/SpeechRecognition), but due to the limitation on TorchScript support, it is constrained on the fixed number of time frames. We want to improve this by adding a model definition that supports scripting method. ## So is this inference only? At the beginning, our focus and the intended use is inference. However, since the model is authored with `torch.nn.Module`, you can still plug-in it to your training as well. Note that the initial version will behave differently in training mode. (Things like random layer drop in Transformer or weight normalization will not be included.) For training, we recommend using `fairseq`, and Hugging Face `transformers`, which are confirmed to work. ## Is there a prototype? You can see the example of running wav2vec2.0 pre-trained models in [the prototype branch](https://github.com/pytorch/audio/tree/wav2vec2-prototype/examples/libtorchaudio/speech_recognition). ## What are the accuracies? With the prototype above, we ran models in C++ and evaluated them on Librispeech. For [models from `fairseq` repository](https://github.com/pytorch/fairseq/tree/master/examples/wav2vec), we also ran them in Python, and got the exact same results. | Architecture | Fine Tune | test-clean | test-other | | ----------------------------------------- | ----------:|-----------:| ----------:| | Base <code>wav2vec_small_960h</code> | 960h | 3.1 | 7.7 | | Large <code>wav2vec_big_960h</code> | 960h | 2.6 | 5.9 | | Large (LV-60) <code>wav2vec2_vox_960h_new</code> | 960h | 2.9 | 6.2 | | Large (LV-60) + Self Training <code>(wav2vec_vox_960h_pl</code> | 960h | 1.9 | 4.5 | For models from Hugging Face `transformers`, we got a similar result. | Architecture | Fine Tune | test-clean | test-other | | ----------------------------------------- | ----------:|-----------:| ----------:| | Base <code>facebook/wav2vec2-base-960h</code> | 960h | 3.1 | 7.8 | | Large <code>facebook/wav2vec2-large-960h</code> | 960h | 2.5 | 5.8 | | Large (LV-60) <code>facebook/wav2vec2-large-960h-lv60</code> | 960h | 3.0 | 6.3 | | Large (LV-60) + Self Training <code>facebook/wav2vec2-large-960h-lv60-self</code> | 960h | 1.9 | 4.5 | We also verified that models fine-tuned using Hugging Face `transformers` work. We ran `facebook/wav2vec2-large-xlsr-53-german` model on VoxForge Germany dataset and got WER of around 10.8. ## What are the performances? We compared the run time performaces of the original `fairseq` model and our C++ version. C++ version is slightly faster but, not significantly so. Quantization of the model might help this. ![Encoding RTF](https://user-images.githubusercontent.com/855818/118184401-4cbb9800-b409-11eb-8180-bf56cfa15211.png) ## What are the future works? There are couple of things we can think of for the future direction. (But please note that we do not promise any future outcome.) 1. Addition of TorchScript-able CTC decoder in `torchaudio`. 2. Quantization support and mobile optimization support. 3. Training support.

If this interests you, please give us a feedback.

Thanks!

Topic		Replies	Views
PyTorch 2.2.0 General Availability Release Announcements	0	1348	January 31, 2024
PyTorch 2.0 Manifesto and Architecture docs compiler	4	1467	December 11, 2022
Torch.nn H2 2021 Lookback and H1 2022 Lookahead frontend API	0	917	January 14, 2022
2.0 \| Feature Review Complete Release Announcements	1	1486	February 10, 2023
Summary of Issues Found in Adapting Models to use TorchScript compiler	1	843	March 16, 2021

Wav2vec2.0 model support in torchaudio

Related topics