add models

README.md

@@ -1,3 +1,84 @@
 ---
-license: apache-2.0
+language: "en"
+thumbnail:
+tags:
+- embeddings
+- Speaker
+- Verification
+- Identification
+- pytorch
+- xvectors
+- TDNN
+- speechbrain
+- audio-classification
+license: "apache-2.0"
+datasets:
+- voxceleb
+inference: false
 ---
+
+# Xvector embeddings extraction on Voxceleb
+
+This repository provides all the necessary tools to extract speaker embeddings with a pretrained TDNN model with SpeechBrain. 
+The system is trained on Voxceleb 1+ Voxceleb2 training data. 
+
+This repo traces the model shared at: https://huggingface.co/speechbrain/spkrec-xvect-voxceleb
+
+For using this model, SpeechBrain is not required.
+
+To use the model:
+
+```python
+import torchaudio
+import torch
+
+device = "cpu"  # or "cuda" for loading the model to a GPU 
+model = torch.jit.load(f"<this_repo>/model_{device}.pt")
+wavsignal, fs = torchaudio.load('audio.wav')
+embeddings = model(wavsignal).squeeze()
+```
+
+## Warning
+
+This model can only forward 1 wave audio.
+Also, this model was traced for speaker embedding extraction, 
+so any additional task (Speaker Classification, Verification) is disabled.
+
+You can retrace the model using the code shared at: "tracing_code.py"
+
+## All the credits to the SpeechBrain Team.
+
+### Limitations
+The SpeechBrain team does not provide any warranty on the performance achieved by this model when used on other datasets.
+
+#### Referencing xvectors
+```@inproceedings{DBLP:conf/odyssey/SnyderGMSPK18,
+  author    = {David Snyder and
+               Daniel Garcia{-}Romero and
+               Alan McCree and
+               Gregory Sell and
+               Daniel Povey and
+               Sanjeev Khudanpur},
+  title     = {Spoken Language Recognition using X-vectors},
+  booktitle = {Odyssey 2018},
+  pages     = {105--111},
+  year      = {2018},
+}
+```
+
+
+# **Citing SpeechBrain**
+Please, cite SpeechBrain if you use it for your research or business.
+
+
+```bibtex
+@misc{speechbrain,
+  title={{SpeechBrain}: A General-Purpose Speech Toolkit},
+  author={Mirco Ravanelli and Titouan Parcollet and Peter Plantinga and Aku Rouhe and Samuele Cornell and Loren Lugosch and Cem Subakan and Nauman Dawalatabad and Abdelwahab Heba and Jianyuan Zhong and Ju-Chieh Chou and Sung-Lin Yeh and Szu-Wei Fu and Chien-Feng Liao and Elena Rastorgueva and François Grondin and William Aris and Hwidong Na and Yan Gao and Renato De Mori and Yoshua Bengio},
+  year={2021},
+  eprint={2106.04624},
+  archivePrefix={arXiv},
+  primaryClass={eess.AS},
+  note={arXiv:2106.04624}
+}
+```

model_cpu.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1acfffabc6ba019df1f9905fc8637e27052075142688d56700bf937edb7e838
+size 16961987

model_cuda.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07568c618fa9f7b164f29b7e6cb2aa56a89240b5cdbaf7857f8e80fadebaa76a
+size 16962241

tracing_code.py

@@ -0,0 +1,92 @@
+import os
+import numpy as np
+
+import torchaudio
+import torch
+from torch import nn
+
+from speechbrain.lobes.models.Xvector import Xvector
+from speechbrain.lobes.features import Fbank
+from speechbrain.processing.features import InputNormalization
+
+
+class Extractor(nn.Module):
+    model_dict = [
+        "mean_var_norm",
+        "compute_features",
+        "embedding_model",
+        "mean_var_norm_emb",
+    ]
+    def __init__(self, model_path, n_mels=24, device="cpu"):
+        super().__init__()
+        self.device = device
+        self.compute_features = Fbank(n_mels=n_mels)
+        self.mean_var_norm = InputNormalization(norm_type="sentence", std_norm=False)
+        self.embedding_model = Xvector(
+            in_channels = n_mels,
+            activation = torch.nn.LeakyReLU,
+            tdnn_blocks = 5,
+            tdnn_channels = [512, 512, 512, 512, 1500],
+            tdnn_kernel_sizes = [5, 3, 3, 1, 1],
+            tdnn_dilations = [1, 2, 3, 1, 1],
+            lin_neurons = 512,
+        )
+        self.mean_var_norm_emb = InputNormalization(norm_type="global", std_norm=False)
+        for mod_name in self.model_dict:
+            filename = os.path.join(model_path, f"{mod_name}.ckpt")
+            module = getattr(self, mod_name)
+            if os.path.exists(filename):
+                if hasattr(module, "_load"):
+                    print(f"Load: {filename}")
+                    module._load(filename)
+                else:
+                    print(f"Load State Dict: {filename}")
+                    module.load_state_dict(torch.load(filename))
+            module.to(self.device)
+
+    @torch.no_grad()
+    def forward(self, wavs, wav_lens = None, normalize=False):
+        # Manage single waveforms in input
+        if len(wavs.shape) == 1:
+            wavs = wavs.unsqueeze(0)
+
+        # Assign full length if wav_lens is not assigned
+        if wav_lens is None:
+            wav_lens = torch.ones(wavs.shape[0], device=self.device)
+
+        # Storing waveform in the specified device
+        wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device)
+        wavs = wavs.float()
+
+        # Computing features and embeddings
+        feats = self.compute_features(wavs)
+        feats = self.mean_var_norm(feats, wav_lens)
+        embeddings = self.embedding_model(feats, wav_lens)
+        if normalize:
+            embeddings = self.mean_var_norm_emb(
+                embeddings, torch.ones(embeddings.shape[0], device=self.device)
+            )
+        return embeddings
+
+
+MODEL_PATH = "pretrained_models/spkrec-xvect-voxceleb"
+signal, fs = torchaudio.load('audio.wav')
+
+device = "cuda"
+extractor = Extractor(MODEL_PATH, device=device)
+
+for k, p in extractor.named_parameters():
+    p.requires_grad = False
+
+extractor.eval()
+embeddings_x = extractor(signal).cpu().squeeze()
+
+# Tracing
+traced_model = torch.jit.trace(extractor, signal)
+torch.jit.save(traced_model, f"model_{device}.pt")
+embeddings_t = traced_model(signal).squeeze()
+print(embeddings_t)
+
+model = torch.jit.load(f"model_{device}.pt")
+emb_m = model(signal).squeeze()
+print(emb_m)