This repository is the official PyTorch implementation of the Robust Speaker-invariant Clustering (R-Spin) proposed in the NAACL 2024 paper R-Spin: Efficient Speaker and Noise-invariant Representation Learning with Acoustic Pieces (Heng-Jui Chang and James Glass; MIT CSAIL).
Please cite our paper if you find this repository and/or the paper useful.
@inproceedings{chang-glass-2024-r,
title = "{R}-Spin: Efficient Speaker and Noise-invariant Representation Learning with Acoustic Pieces",
author = "Chang, Heng-Jui and Glass, James",
booktitle = "Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)",
month = jun,
year = "2024",
address = "Mexico City, Mexico",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2024.naacl-long.36",
doi = "10.18653/v1/2024.naacl-long.36",
pages = "642--662",
}- Install the latest PyTorch and soundfile.
pip install torch soundfile librosa
- Download checkpoint (WavLM + R-Spin):
| Spin Codebook Size | Acoustic Pieces | Checkpoint |
|---|---|---|
| 32 (best for ASR) | 40k | link |
| 64 | 40k | link |
| 128 | 40k | link |
| 256 | 40k | link |
| 512 | 40k | link |
| 1024 | 40k | link |
| 2048 | 40k | link |
import soundfile as sf
import torch
from rspin import RSpinWavlm
# Load model
model = RSpinWavlm.load_from_checkpoint("/path/to/checkpoint").cuda()
model.eval()
# Load audio (needs to be 16kHz)
wav, sr = sf.read("/path/to/audio")
assert sr == 16000
if wav.ndim == 2:
wav = wav.mean(-1)
wav = torch.FloatTensor(wav).cuda()
# Inference
with torch.inference_mode():
feat_list, padding_mask, codes = model(wav.unsqueeze(0), get_code=True)
print(codes[0])
# feat_list: List[torch.Tensor] (shape: (batch, seq_len, encoder_emb_dim))
# codes: torch.LongTensor (shape: (batch, seq_len))- The
feat_listconsists a list of all hidden representations (including the CNN feature extractor output). - The
codesis aLongTensorrepresenting the codeword IDs produced by the Spin codebook. - See rspin/model.py for more information.
We use the s3prl toolkit for SUPERB downstream tasks.
- Modify line 24 of
s3prl_py/rspin/expert.pyto the absolute path torspin/. - Copy the
s3prl_py/rspindirectory tos3prlso that the toolkit can load the models.
cp -R s3prl_py/rspin ../s3prl/s3prl/upstream/rspin- Add the following line to
../s3prl/s3prl/hub.py:
from s3prl.upstream.rspin.hubconf import *- Try loading R-Spin models with the following methods
from s3prl.nn import S3PRLUpstream
# Method 1 (download checkpoints manually with the provided links)
model = S3PRLUpstream("rspin_local", path_or_url="/path/to/checkpoint")
# Method 2 (download checkpoint with s3prl)
model = S3PRLUpstream("rspin_wavlm_32_40k")After training an ASR model with the s3prl toolkit, you may use the model to transcribe speech.
import torch
from rspin import RSpinASR
device = torch.device("cuda")
asr_model = RSpinASR.load_from_checkpoint(
"/path/to/s3prl/checkpoint", "/path/to/rspin/checkpoint", device=device
)
asr_model.eval()
with torch.no_grad():
audio_paths = ["/audio/file/1", "/audio/file/2"]
results = asr_model(audio_paths=audio_paths)
for transcription in results["transcription"]:
print(transcription)If you have any questions, please open an issue or email hengjui [at] mit.edu.