Fast Pytorch Kmeans

this is a pytorch implementation of K-means clustering algorithm

Quick Start

from kmeans import KMeans
import torch

kmeans = KMeans(n_clusters=8, mode='euclidean', verbose=1)
x = torch.randn(100000, 64, device='cuda')
labels = kmeans.fit_predict(x)

Speed Comparison

Tested on google colab with a Tesla K80 GPU

n_samples=100,000 n_features=64 iterations=100

n_samples=100,000 n_clusters=64 iterations=100

n_features=256, n_clusters=256 iterations=100

n_features=512, n_clusters=512 iterations=100