kmeans_pytorch

kmeans_pytorch copied to clipboard

pytorch implementation of basic kmeans algorithm(lloyd method with forgy initialization) with gpu support

usage:

from kmeans_pytorch.kmeans import lloyd
import numpy as np 

A = np.concatenate([np.random.randn(1000, 2), p.random.randn(1000, 2)+3, p.random.randn(1000, 2)+6], axis=0)

#lloyd(X, n_clusters, device=0, tol=1e-4)
clusters_index, centers = lloyd(A, 2, device=0, tol=1e-4)

See kmeans_test.ipynb for some test examples