Kmeans是一种简单易用的聚类算法,是少有的会出现在深度学习项目中的传统算法,比如人脸搜索项目、物体检测项目(yolov3中用到了Kmeans进行anchors聚类)等。
一般使用Kmeans会直接调sklearn,如果任务比较复杂,可以通过numpy进行自定义,这里介绍使用Pytorch实现的方式,经测试,通过Pytorch调用GPU之后,能够提高多特征聚类的速度。
import torch
import time
from tqdm import tqdm
class KMEANS:
def __init__(self, n_clusters=20, max_iter=None, verbose=True,device = torch.device("cpu")):
self.n_cluster = n_clusters
self.n_clusters = n_clusters
self.labels = None
self.dists = None # shape: [x.shape[0],n_cluster]
self.centers = None
self.variation = torch.Tensor([float("Inf")]).to(device)
self.verbose = verbose
self.started = False
self.representative_samples = None
self.max_iter = max_iter
self.count = 0
self.device =