Add ONNX compatibility (#204)

* Add ONNX segmentation and embedding models

* Minor readability improvements

* Replace onnxruntime with onnxruntime-gpu

* ONNX can have only one output

* Clean up useless embedding model subclasses

* Remove duration and sample_rate properties from SegmentationModel. Clean up code

* Update README

Author: juanmc2005

README.md

@@ -123,7 +123,7 @@ For inference and evaluation on a dataset we recommend to use `Benchmark` (see n
 
 ## 🤖 Add your model
 
-Third-party models can be integrated by subclassing `SegmentationModel` and `EmbeddingModel` (both PyTorch `nn.Module`):
+Third-party models can be integrated by providing a loader function:
 
 ```python
 from diart import SpeakerDiarization, SpeakerDiarizationConfig
@@ -132,46 +132,39 @@ from diart.sources import MicrophoneAudioSource
 from diart.inference import StreamingInference
 
 
-def model_loader():
+def segmentation_loader():
+    # It should take a waveform and return a segmentation tensor
     return load_pretrained_model("my_model.ckpt")
 
+def embedding_loader():
+    # It should take (waveform, weights) and return per-speaker embeddings
+    return load_pretrained_model("my_other_model.ckpt")
 
-class MySegmentationModel(SegmentationModel):
-    def __init__(self):
-        super().__init__(model_loader)
-    
-    @property
-    def sample_rate(self) -> int:
-        return 16000
-    
-    @property
-    def duration(self) -> float:
-        return 2  # seconds
-    
-    def forward(self, waveform):
-        # self.model is created lazily
-        return self.model(waveform)
-
-    
-class MyEmbeddingModel(EmbeddingModel):
-    def __init__(self):
-        super().__init__(model_loader)
-    
-    def forward(self, waveform, weights):
-        # self.model is created lazily
-        return self.model(waveform, weights)
-
-    
+
+segmentation = SegmentationModel(segmentation_loader)
+embedding = EmbeddingModel(embedding_loader)
 config = SpeakerDiarizationConfig(
-    segmentation=MySegmentationModel(),
-    embedding=MyEmbeddingModel()
+    segmentation=segmentation,
+    embedding=embedding,
 )
 pipeline = SpeakerDiarization(config)
 mic = MicrophoneAudioSource()
 inference = StreamingInference(pipeline, mic)
 prediction = inference()
 ```
 
+If you have an ONNX model, you can use `from_onnx()`:
+
+```python
+from diart.models import EmbeddingModel
+
+embedding = EmbeddingModel.from_onnx(
+    model_path="my_model.ckpt",
+    input_names=["x", "w"],  # defaults to ["waveform", "weights"]
+    output_name="output",  # defaults to "embedding"
+)
+```
+
 ## 📈 Tune hyper-parameters
 
 Diart implements an optimizer based on [optuna](https://optuna.readthedocs.io/en/stable/index.html) that allows you to tune pipeline hyper-parameters to your needs.
@@ -352,11 +345,11 @@ from diart.models import SegmentationModel
 
 benchmark = Benchmark("/wav/dir", "/rttm/dir")
 
-name = "pyannote/segmentation@Interspeech2021"
-segmentation = SegmentationModel.from_pyannote(name)
+model_name = "pyannote/segmentation@Interspeech2021"
+model = SegmentationModel.from_pretrained(model_name)
 config = SpeakerDiarizationConfig(
-    # Set the model used in the paper
-    segmentation=segmentation,
+    # Set the segmentation model used in the paper
+    segmentation=model,
     step=0.5,
     latency=0.5,
     tau_active=0.555,

environment.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - python=3.8
+  - python=3.10
   - portaudio=19.6.*
   - pysoundfile=0.12.*
   - ffmpeg[version='<4.4']

requirements.txt

@@ -18,3 +18,4 @@ optuna>=2.10
 websocket-server>=0.6.4
 websocket-client>=0.58.0
 rich>=12.5.1
+onnxruntime-gpu>=1.16.1

setup.cfg

@@ -40,6 +40,7 @@ install_requires=
     websocket-server>=0.6.4
     websocket-client>=0.58.0
     rich>=12.5.1
+    onnxruntime-gpu>=1.16.1
 
 [options.packages.find]
 where=src

src/diart/blocks/diarization.py

@@ -23,7 +23,7 @@ def __init__(
         self,
         segmentation: m.SegmentationModel | None = None,
         embedding: m.EmbeddingModel | None = None,
-        duration: float | None = None,
+        duration: float = 5,
         step: float = 0.5,
         latency: float | Literal["max", "min"] | None = None,
         tau_active: float = 0.6,
@@ -34,6 +34,7 @@ def __init__(
         max_speakers: int = 20,
         normalize_embedding_weights: bool = False,
         device: torch.device | None = None,
+        sample_rate: int = 16000,
         **kwargs,
     ):
         # Default segmentation model is pyannote/segmentation
@@ -47,7 +48,7 @@ def __init__(
         )
 
         self._duration = duration
-        self._sample_rate: int | None = None
+        self._sample_rate = sample_rate
 
         # Latency defaults to the step duration
         self._step = step
@@ -70,9 +71,6 @@ def __init__(
 
     @property
     def duration(self) -> float:
-        # Default duration is the one given by the segmentation model
-        if self._duration is None:
-            self._duration = self.segmentation.duration
         return self._duration
 
     @property
@@ -85,9 +83,6 @@ def latency(self) -> float:
 
     @property
     def sample_rate(self) -> int:
-        # Expected sample rate is given by the segmentation model
-        if self._sample_rate is None:
-            self._sample_rate = self.segmentation.sample_rate
         return self._sample_rate
 
 
@@ -177,9 +172,8 @@ def __call__(
 
         # Extract segmentation and embeddings
         segmentations = self.segmentation(batch)  # shape (batch, frames, speakers)
-        embeddings = self.embedding(
-            batch, segmentations
-        )  # shape (batch, speakers, emb_dim)
+        # embeddings has shape (batch, speakers, emb_dim)
+        embeddings = self.embedding(batch, segmentations)
 
         seg_resolution = waveforms[0].extent.duration / segmentations.shape[1]
 

src/diart/blocks/embedding.py

@@ -3,6 +3,7 @@
 import torch
 from einops import rearrange
 
+from .. import functional as F
 from ..features import TemporalFeatures, TemporalFeatureFormatter
 from ..models import EmbeddingModel
 
@@ -90,10 +91,8 @@ def __init__(self, gamma: float = 3, beta: float = 10, normalize: bool = False):
 
     def __call__(self, segmentation: TemporalFeatures) -> TemporalFeatures:
         weights = self.formatter.cast(segmentation)  # shape (batch, frames, speakers)
-        with torch.no_grad():
-            probs = torch.softmax(self.beta * weights, dim=-1)
-            weights = torch.pow(weights, self.gamma) * torch.pow(probs, self.gamma)
-            weights[weights < 1e-8] = 1e-8
+        with torch.inference_mode():
+            weights = F.overlapped_speech_penalty(weights, self.gamma, self.beta)
             if self.normalize:
                 min_values = weights.min(dim=1, keepdim=True).values
                 max_values = weights.max(dim=1, keepdim=True).values
@@ -110,19 +109,8 @@ def __init__(self, norm: Union[float, torch.Tensor] = 1):
             self.norm = self.norm.unsqueeze(0)
 
     def __call__(self, embeddings: torch.Tensor) -> torch.Tensor:
-        # Add batch dimension if missing
-        if embeddings.ndim == 2:
-            embeddings = embeddings.unsqueeze(0)
-        if isinstance(self.norm, torch.Tensor):
-            batch_size1, num_speakers1, _ = self.norm.shape
-            batch_size2, num_speakers2, _ = embeddings.shape
-            assert batch_size1 == batch_size2 and num_speakers1 == num_speakers2
-        with torch.no_grad():
-            norm_embs = (
-                self.norm
-                * embeddings
-                / torch.norm(embeddings, p=2, dim=-1, keepdim=True)
-            )
+        with torch.inference_mode():
+            norm_embs = F.normalize_embeddings(embeddings, self.norm)
         return norm_embs
 
 

src/diart/blocks/vad.py

@@ -27,11 +27,12 @@ class VoiceActivityDetectionConfig(base.PipelineConfig):
     def __init__(
         self,
         segmentation: m.SegmentationModel | None = None,
-        duration: float | None = None,
+        duration: float = 5,
         step: float = 0.5,
         latency: float | Literal["max", "min"] | None = None,
         tau_active: float = 0.6,
         device: torch.device | None = None,
+        sample_rate: int = 16000,
         **kwargs,
     ):
         # Default segmentation model is pyannote/segmentation
@@ -41,7 +42,7 @@ def __init__(
 
         self._duration = duration
         self._step = step
-        self._sample_rate: int | None = None
+        self._sample_rate = sample_rate
 
         # Latency defaults to the step duration
         self._latency = latency
@@ -57,9 +58,6 @@ def __init__(
 
     @property
     def duration(self) -> float:
-        # Default duration is the one given by the segmentation model
-        if self._duration is None:
-            self._duration = self.segmentation.duration
         return self._duration
 
     @property
@@ -72,9 +70,6 @@ def latency(self) -> float:
 
     @property
     def sample_rate(self) -> int:
-        # Expected sample rate is given by the segmentation model
-        if self._sample_rate is None:
-            self._sample_rate = self.segmentation.sample_rate
         return self._sample_rate
 
 

src/diart/console/benchmark.py

@@ -43,6 +43,7 @@ def run():
     parser.add_argument(
         "--duration",
         type=float,
+        default=5,
         help=f"{argdoc.DURATION}. Defaults to training segmentation duration",
     )
     parser.add_argument(
@@ -111,8 +112,8 @@ def run():
 
     # Resolve models
     hf_token = utils.parse_hf_token_arg(args.hf_token)
-    args.segmentation = m.SegmentationModel.from_pyannote(args.segmentation, hf_token)
-    args.embedding = m.EmbeddingModel.from_pyannote(args.embedding, hf_token)
+    args.segmentation = m.SegmentationModel.from_pretrained(args.segmentation, hf_token)
+    args.embedding = m.EmbeddingModel.from_pretrained(args.embedding, hf_token)
 
     pipeline_class = utils.get_pipeline_class(args.pipeline)
 

src/diart/console/serve.py

@@ -36,6 +36,7 @@ def run():
     parser.add_argument(
         "--duration",
         type=float,
+        default=5,
         help=f"{argdoc.DURATION}. Defaults to training segmentation duration",
     )
     parser.add_argument(
@@ -92,8 +93,8 @@ def run():
 
     # Resolve models
     hf_token = utils.parse_hf_token_arg(args.hf_token)
-    args.segmentation = m.SegmentationModel.from_pyannote(args.segmentation, hf_token)
-    args.embedding = m.EmbeddingModel.from_pyannote(args.embedding, hf_token)
+    args.segmentation = m.SegmentationModel.from_pretrained(args.segmentation, hf_token)
+    args.embedding = m.EmbeddingModel.from_pretrained(args.embedding, hf_token)
 
     # Resolve pipeline
     pipeline_class = utils.get_pipeline_class(args.pipeline)

src/diart/console/stream.py

@@ -39,6 +39,7 @@ def run():
     parser.add_argument(
         "--duration",
         type=float,
+        default=5,
         help=f"{argdoc.DURATION}. Defaults to training segmentation duration",
     )
     parser.add_argument(
@@ -103,8 +104,8 @@ def run():
 
     # Resolve models
     hf_token = utils.parse_hf_token_arg(args.hf_token)
-    args.segmentation = m.SegmentationModel.from_pyannote(args.segmentation, hf_token)
-    args.embedding = m.EmbeddingModel.from_pyannote(args.embedding, hf_token)
+    args.segmentation = m.SegmentationModel.from_pretrained(args.segmentation, hf_token)
+    args.embedding = m.EmbeddingModel.from_pretrained(args.embedding, hf_token)
 
     # Resolve pipeline
     pipeline_class = utils.get_pipeline_class(args.pipeline)

src/diart/console/tune.py

@@ -46,6 +46,7 @@ def run():
     parser.add_argument(
         "--duration",
         type=float,
+        default=5,
         help=f"{argdoc.DURATION}. Defaults to training segmentation duration",
     )
     parser.add_argument(
@@ -120,8 +121,8 @@ def run():
 
     # Resolve models
     hf_token = utils.parse_hf_token_arg(args.hf_token)
-    args.segmentation = m.SegmentationModel.from_pyannote(args.segmentation, hf_token)
-    args.embedding = m.EmbeddingModel.from_pyannote(args.embedding, hf_token)
+    args.segmentation = m.SegmentationModel.from_pretrained(args.segmentation, hf_token)
+    args.embedding = m.EmbeddingModel.from_pretrained(args.embedding, hf_token)
 
     # Retrieve pipeline class
     pipeline_class = utils.get_pipeline_class(args.pipeline)

src/diart/functional.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import torch
+
+
+def overlapped_speech_penalty(
+    segmentation: torch.Tensor, gamma: float = 3, beta: float = 10
+):
+    # segmentation has shape (batch, frames, speakers)
+    probs = torch.softmax(beta * segmentation, dim=-1)
+    weights = torch.pow(segmentation, gamma) * torch.pow(probs, gamma)
+    weights[weights < 1e-8] = 1e-8
+    return weights
+
+
+def normalize_embeddings(
+    embeddings: torch.Tensor, norm: float | torch.Tensor = 1
+) -> torch.Tensor:
+    # embeddings has shape (batch, speakers, feat) or (speakers, feat)
+    if embeddings.ndim == 2:
+        embeddings = embeddings.unsqueeze(0)
+    if isinstance(norm, torch.Tensor):
+        batch_size1, num_speakers1, _ = norm.shape
+        batch_size2, num_speakers2, _ = embeddings.shape
+        assert batch_size1 == batch_size2 and num_speakers1 == num_speakers2
+    emb_norm = torch.norm(embeddings, p=2, dim=-1, keepdim=True)
+    return norm * embeddings / emb_norm