带有gpu支持的基本kmeans算法(带有Forgy初始化的劳埃德方法)的pytorch实现
用法:
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 )
有关一些测试示例,请参见kmeans_test.ipynb
1