feat(clustering): Add KMeans GPU clustering support

pyannote/audio/pipelines/clustering.py

@@ -471,6 +471,79 @@ def cluster(
  return clusters
 
 
+class KMeansGPU(BaseClustering):
+ def __init__(
+ self,
+ metric: str = "",
+ max_num_embeddings: int = np.inf,
+ constrained_assignment: bool = False,
+ ):
+ """KMeans clustering
+
+ Parameters
+ ----------
+ metric : {""}, optional
+ Distance metric to use. KMeansGPU only supports the default value.
+ """
+ super().__init__(
+ metric=metric,
+ max_num_embeddings=max_num_embeddings,
+ constrained_assignment=constrained_assignment,
+ )
+
+ def cluster(
+ self, embeddings, min_clusters: int, max_clusters: int, num_clusters: int = None
+ ):
+ try:
+ import cuml
+ import cupy as cp
+ from cuml.metrics.cluster import silhouette_score
+ except ImportError:
+ raise ImportError(
+ "KMeansGPU requires cuML. You can install it with 'https://docs.rapids.ai/install'."
+ )
+
+ assert max_clusters >= min_clusters > 0
+
+ num_embeddings = len(embeddings)
+
+ may_single = False
+ if max_clusters > 1 and min_clusters == 1:
+ min_clusters = 2
+ may_single = True
+ elif max_clusters == 1:
+ return np.zeros((num_embeddings,))
+
+ if num_embeddings <= min_clusters or num_embeddings == num_clusters:
+ return np.arange(num_embeddings)
+
+ if num_clusters is not None:
+ agg_clust = cuml.cluster.KMeans(n_clusters=num_clusters)
+ clusters = agg_clust.fit_predict(embeddings)
+ return clusters.get()
+
+ embeddings = cp.asarray(embeddings)
+
+ best_score = -1
+ best_clusters = None
+
+ for num_clusters in range(min_clusters, min(max_clusters + 1, num_embeddings)):
+ agg_clust = cuml.cluster.KMeans(n_clusters=num_clusters)
+ clusters = agg_clust.fit_predict(embeddings)
+
+ score = silhouette_score(embeddings, clusters)
+
+ if score > best_score:
+ best_score = score
+ best_clusters = clusters
+
+ if may_single:
+ if num_clusters == 2 and best_score < 0.25:
+ return np.zeros((num_embeddings,))
+
+ return best_clusters.get()
+
+
 class OracleClustering(BaseClustering):
  """Oracle clustering"""
 
@@ -558,4 +631,5 @@ def __call__(
 
 class Clustering(Enum):
  AgglomerativeClustering = AgglomerativeClustering
+ KMeansGPU = KMeansGPU
  OracleClustering = OracleClustering

tests/test_clustering.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pyannote.audio.pipelines.clustering import AgglomerativeClustering
+from pyannote.audio.pipelines.clustering import AgglomerativeClustering, KMeansGPU
 
 
 def test_agglomerative_clustering_num_cluster():
@@ -26,4 +26,83 @@ def test_agglomerative_clustering_num_cluster():
  clusters = clustering.cluster(
  embeddings=embeddings, min_clusters=2, max_clusters=2, num_clusters=2
  )
+ print(clusters)
  assert np.array_equal(clusters, np.array([0, 1]))
+
+
+def test_kmeans_clustering_num_cluster_gpu_too_small():
+ clustering = KMeansGPU().instantiate({})
+
+ embeddings = np.array([[1.0, 1.0, 1.0, 1.0], [1.0, 2.0, 1.0, 2.0]])
+
+ # request 2 clusters
+ clusters = clustering.cluster(
+ embeddings=embeddings, min_clusters=2, max_clusters=2, num_clusters=2
+ )
+
+ assert np.array_equal(clusters, np.array([0, 1]))
+
+ # generate a 256-dimensional random vector
+ v = np.random.rand(256)
+
+ # define the range and standard deviation of the generated cluster center
+ cluster_center_std_dev = 2.0
+
+ # generate 8 cluster centers randomly
+ num_clusters = 8
+ cluster_centers = np.random.normal(
+ np.mean(v), cluster_center_std_dev, size=(num_clusters, 256)
+ )
+
+ a, b, c = clustering.set_num_clusters(num_clusters, 10, 1, 10)
+ assert a == 8
+
+ a, b, c = clustering.set_num_clusters(num_clusters, None, 1, 10)
+ assert a is None and b == 1 and c == 8
+
+ a, b, c = clustering.set_num_clusters(num_clusters, None, 8, 10)
+ assert a == 8
+
+ a, b, c = clustering.set_num_clusters(num_clusters, None, 7, 10)
+ assert a is None and b == 7 and c == 8
+
+ clustering.cluster(
+ embeddings=cluster_centers, num_clusters=a, min_clusters=b, max_clusters=c
+ )
+
+
+def test_kmeans_clustering_num_cluster_gpu_large():
+ clustering = KMeansGPU().instantiate({})
+
+ # generate a 256-dimensional random vector
+ v = np.random.rand(256)
+
+ # define the range and standard deviation of the generated cluster center
+ cluster_center_std_dev = 2.0
+ vector_std_dev = 1
+
+ # generate 5 cluster centers randomly
+ num_clusters = 5
+ cluster_centers = np.random.normal(
+ np.mean(v), cluster_center_std_dev, size=(num_clusters, 256)
+ )
+
+ # generate 2000 * 32 vectors
+ num_vectors_per_cluster = int(2000 * 32 / num_clusters)
+ all_vectors = []
+
+ for center in cluster_centers:
+ vectors = np.random.normal(
+ center, vector_std_dev, size=(num_vectors_per_cluster, 256)
+ )
+ all_vectors.append(vectors)
+
+ # stack all vectors
+ all_vectors = np.vstack(all_vectors)
+
+ np.random.shuffle(all_vectors)
+
+ clusters = clustering.cluster(
+ embeddings=all_vectors, min_clusters=2, max_clusters=10
+ )
+ assert np.unique(clusters).shape[0] == num_clusters