Klitter / kmeans_pytorch

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

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

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

About

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


Languages

Language:Python 100.0%