fix(diarization): scale clustering with known speaker counts

funasr/models/campplus/cluster_backend.py

@@ -11,6 +11,7 @@ import numpy as np
 
 from sklearn.cluster._kmeans import k_means
 from sklearn.cluster import HDBSCAN
+from sklearn.preprocessing import normalize
 
 
 class SpectralCluster:
@@ -191,6 +192,21 @@ class UmapHdbscan:
         return labels
 
 
+class KMeansCluster:
+    r"""Cluster a large set of speaker embeddings into a known number of groups."""
+
+    def __call__(self, X, num_clusters):
+        """Cluster L2-normalized embeddings with bounded memory usage."""
+        normalized_X = normalize(X)
+        _, labels, _ = k_means(
+            normalized_X,
+            num_clusters,
+            random_state=0,
+            n_init=10,
+        )
+        return labels
+
+
 class ClusterBackend(torch.nn.Module):
     r"""Perfom clustering for input embeddings and output the labels.
     Args:
@@ -210,6 +226,7 @@ class ClusterBackend(torch.nn.Module):
 
         self.spectral_cluster = SpectralCluster()
         self.umap_hdbscan_cluster = UmapHdbscan()
+        self.kmeans_cluster = KMeansCluster()
 
     def forward(self, X, **params):
         # clustering and return the labels
@@ -223,9 +240,10 @@ class ClusterBackend(torch.nn.Module):
         assert len(X.shape) == 2, "modelscope error: the shape of input should be [N, C]"
         if X.shape[0] < 20:
             return np.zeros(X.shape[0], dtype="int")
-        if X.shape[0] < 2048 or k is not None:
-            # unexpected corner case
+        if X.shape[0] < 2048:
             labels = self.spectral_cluster(X, k)
+        elif k is not None:
+            labels = self.kmeans_cluster(X, k)
         else:
             labels = self.umap_hdbscan_cluster(X)
 

tests/test_cluster_backend.py

@@ -0,0 +1,44 @@
+import numpy as np
+import pytest
+import torch
+
+from funasr.models.campplus.cluster_backend import ClusterBackend
+
+
+def test_large_known_speaker_count_uses_fixed_k_clustering(monkeypatch):
+    backend = ClusterBackend()
+    embeddings = torch.ones((2048, 2))
+    expected = np.arange(embeddings.shape[0]) % 2
+
+    def fail_spectral(*args, **kwargs):
+        pytest.fail("large inputs must not use dense spectral clustering")
+
+    def fail_umap(*args, **kwargs):
+        pytest.fail("a known speaker count should use fixed-K clustering")
+
+    def fixed_k_cluster(actual_embeddings, num_clusters):
+        assert actual_embeddings is embeddings
+        assert num_clusters == 2
+        return expected
+
+    monkeypatch.setattr(backend, "spectral_cluster", fail_spectral)
+    monkeypatch.setattr(backend, "umap_hdbscan_cluster", fail_umap)
+    monkeypatch.setattr(backend, "kmeans_cluster", fixed_k_cluster, raising=False)
+
+    labels = backend(embeddings, oracle_num=2)
+
+    np.testing.assert_array_equal(labels, expected)
+
+
+def test_large_fixed_k_clustering_separates_cosine_clusters():
+    backend = ClusterBackend()
+    embeddings = torch.zeros((2048, 2))
+    embeddings[:1024, 0] = 1
+    embeddings[1024:, 1] = 1
+
+    labels = backend(embeddings, oracle_num=2)
+
+    assert np.unique(labels).size == 2
+    assert np.unique(labels[:1024]).size == 1
+    assert np.unique(labels[1024:]).size == 1
+    assert labels[0] != labels[-1]