Benchmark

K-Means

Performing K-Means clustering on float32 data randomly sampled from normal distribution.

Number of iterations is set to 15.
Tolerance is set to 0 in order to perform full 15 iterations of K-Means
Initial centroids are randomly chosen from training data
All runs are performed on a Tesla T4 GPU

Contestants:

torchpq.clustering.KMeans
faiss.Clustering
KeOps

n_features=256, n_clusters=256, varying n_data

n_features=256, n_clusters=16384, varying n_data

n_features=128, n_data=1,000,000, varying n_clusters

n_clusters=1024, n_data=1,000,000, varying n_features

note: faiss and keOps went OOM when n_features > 512