SEW (Squeezed and Efficient Wav2vec)
The repo contains the code of the paper "Performance-Efficiency Trade-offs in Unsupervised Pre-training for Speech Recognition" by Felix Wu, Kwangyoun Kim, Jing Pan, Kyu Han, Kilian Q Weinberger, and Yoav Artzi.
Model Checkpoints
Unsupervisedly Pre-trained on LibriSpeech 960h
| Model | Pre-training updates | Dataset | Model |
|---|---|---|---|
| W2V2-tiny | 100K | Librispeech 960h | download |
| W2V2-small | 100K | Librispeech 960h | download |
| W2V2-mid | 100K | Librispeech 960h | download |
| W2V2-base | 100K | Librispeech 960h | download |
| SEW-tiny | 100K | Librispeech 960h | download |
| SEW-small | 100K | Librispeech 960h | download |
| SEW-mid | 100K | Librispeech 960h | download |
| SEW-D-tiny | 100K | Librispeech 960h | download |
| SEW-D-small | 100K | Librispeech 960h | download |
| SEW-D-mid | 100K | Librispeech 960h | download |
| SEW-D-mid (k127) | 100K | Librispeech 960h | download |
| SEW-D-base | 100K | Librispeech 960h | download |
| SEW-D-base+ | 100K | Librispeech 960h | download |
| SEW-D-mid | 400K | Librispeech 960h | download |
| SEW-D-mid (k127) | 400K | Librispeech 960h | download |
| SEW-D-base+ | 400K | Librispeech 960h | download |
Usage
Dependencies
The code is tested with fairseq commit 05255f9, deberta commit bf17ca4 and the following packages.
torch==1.8.0
torchaudio==0.8.0
tqdm==4.49.0
Hydra==2.5
hydra-core==1.0.4
fvcore==0.1.5.post20210330
omegaconf==2.0.5
einops==0.3.0
fire==0.2.1
Apex
Please install NVIDIA's apex with
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" \
--global-option="--deprecated_fused_adam" --global-option="--xentropy" \
--global-option="--fast_multihead_attn" ./
wav2letter decoder
Currently, we are decoding with wav2letter v0.2 python binding at commit 96f5f9d Please install the python binding here https://github.com/flashlight/wav2letter/tree/96f5f9d3b41e01af0a031ee0d2604acd9ef3b1b0/bindings/python The newest commit d5a93f0 in v0.2 branch leads to worse WER for wav2vec 2.0 baselines.
Installation
git clone https://github.com/asappresearch/sew.git
cd sew
pip install -e .
Pre-training
Pre-training SEW models
Run the following command where $model_size can be tiny, small, or mid, and $ngpu is tne number of GPUs you want to use.
bash scripts/pt-sew.sh $model_size $ngpu
Pre-training SEW-D models
bash scripts/pt-sew-d.sh $model_size $ngpu
where $model_size can be tiny, small, mid, mid-k127, base, or base+.
Fine-tuning
Run the following script to fine-tune a model with the hyperparameters from wav2vec 2.0.
bash scripts/ft-model.sh $pre_trained_model $split $ngpu
where $pre_trained_model can be either a W2V2, SEW, or a SEW-D model checkpoint and $split can be 10m, 1h, 10h, or 100h.
Here we also provide a set of hyperparameters which sets all dropouts the same as the pre-training stage, and we found it to be more stable.
bash scripts/ft-model-stable.sh $pre_trained_model $split $ngpu
If you see out of GPU memory error, please scale down the dataset.max_tokens and scale up the optimization.update_freq in scripts/ft-model.sh. For example modifying these lines
dataset.max_tokens=3200000 \
optimization.update_freq="[$((8 / $ngpu))]" \
to
dataset.max_tokens=1600000 \
optimization.update_freq="[$((16 / $ngpu))]" \
which reduces the batch size and increases the gradient accumulation steps in order to use less GPU memory.
Evaluation
- Please run this script to prepare the official LibriSpeech 4-gram language model.
bash scripts/prepare_librispeech_lm.sh $kenlm_build_bin
where $kenlm_build_bin is the folder that contains the KenLM build_binary executable file (e.g. /home/user/kenlm/build/bin).
- Then run this script to evaluate a pre-trained ASR model
python tools/eval_w2v.py tunelm --subsets '["dev-clean", "dev-other", "test-clean", "test-other"]' --model $asr_checkpoint
139 Nov 18, 2022
43 Nov 19, 2022
37 Nov 21, 2022
199 Jan 8, 2023
296 Dec 29, 2022
138 Apr 15, 2022
56 Nov 28, 2022
59 Oct 31, 2022
1 Nov 04, 2020
272 Jan 08, 2023
170 Dec 15, 2022
333 Dec 29, 2022
1.4k Jan 08, 2023
106 Jan 09, 2023
13 May 11, 2022
6 Jan 18, 2022
1 Dec 10, 2021
53 Dec 13, 2022
6 Dec 21, 2022
9 Dec 04, 2022
61 Dec 25, 2022
2.5k Jan 04, 2023
36 Jan 01, 2023
16 Jul 16, 2022
1k Jan 03, 2023
60 Dec 22, 2022
36 Dec 16, 2022
0 Jan 08, 2022