Squeezeformer: An Efficient Transformer for Automatic Speech Recognition

We provide testing codes for Squeezeformer, along with the pre-trained checkpoints.

Check out paper for more details.

Install Squeezeformer

We recommend using Python version 3.8.

1. Install dependancies

We support Tensorflow version of 2.5. Run the following commands depending on your target device type.

Running on CPUs: pip install -e '.[tf2.5]'
Running on GPUs: pip install -e '.[tf2.5-gpu]'

2. Install CTC decoder

cd scripts
bash install_ctc_decoders.sh

Prepare Dataset

1. Download Librispeech

Librispeech is a widely-used ASR benchmark that consists of 960hr speech corpus with text transcriptions. The dataset consists of 3 training sets (train-clean-100, train-clean-360, train-other-500), 2 development sets (dev-clean, dev-other), and 2 test sets (test-clean, test-other).

Download the datasets from this link and untar them. If this is for testing purposes only, you can skip the training datasets to save disk space. You should have flac files under {dataset_path}/LibriSpeech.

2. Create Manifest Files

Once you download the datasets, you should create a manifest file that links the file path to the audio input and its transcription. We use a script from TensorFlowASR.

cd scripts
python create_librispeech_trans_all.py --data {dataset_path}/LibriSpeech --output {tsv_dir}

The dataset_path is the directory that you untarred the datasets in the previous step.
This script creates tsv files under tsv_dir that list the audio file path, duration, and the transcription.
To skip processing the training datasets, use an additional argument --mode test-only.

If you have followed the instruction correctly, you should have the following files under tsv_dir.

dev_clean.tsv, dev_other.tsv, test_clean.tsv, test_other.tsv
train_clean_100.tsv, train_clean_360.tsv, train_other_500.tsv (if not --mode test-only)
train_other.tsv that merges all training tsv files into one (if not --mode test-only)

Testing Squeezeformer

1. Download Pre-trained Checkpoints

We provide pre-trained checkpoints for all variants of Squeezeformer.

Model	Checkpoint	test-clean	test-other
Squeezeformer-XS	link	3.74	9.09
Squeezeformer-S	link	3.08	7.47
Squeezeformer-SM	link	2.79	6.89
Squeezeformer-M	link	2.56	6.50
Squeezeformer-ML	link	2.61	6.05
Squeezeformer-L	link	2.47	5.97

2. Run Inference!

Run the following commands:

cd examples/squeezeformer
python test.py --bs {batch_size} --config configs/squeezeformer-S.yml --saved squeezeformer-S.h5 \
    --dataset_path {tsv_dir} --dataset {dev_clean|dev_other|test_clean|test_other}

tsv_dir is the directory path to the tsv manifest files that you created in the previous step.
You can test on other Squeezeformer models by changing --config and --saved, e.g., Squeezeformer-L or Squeezeformer-M.

External implementations

We are thankful to all the researchers who have extended Squeezeformer for different purposes.

Description	Checkpoint
PyTorch implementation	link

Citation

Squeezeformer has been developed as part of the following paper. We appreciate it if you would please cite the following paper if you found the library useful for your work:

@article{kim2022squeezeformer,
  title={Squeezeformer: An Efficient Transformer for Automatic Speech Recognition},
  author={Kim, Sehoon and Gholami, Amir and Shaw, Albert and Lee, Nicholas and Mangalam, Karttikeya and Malik, Jitendra and Mahoney, Michael W and Keutzer, Kurt},
  journal={arxiv:2206.00888},
  year={2022}
}

kobenaxie/Squeezeformer