Kmeans是一種簡(jiǎn)單易用的聚類算法，是少有的會(huì)出現(xiàn)在深度學(xué)習(xí)項(xiàng)目中的傳統(tǒng)算法，比如人臉?biāo)阉黜?xiàng)目、物體檢測(cè)項(xiàng)目（yolov3中用到了Kmeans進(jìn)行anchors聚類）等。

一般使用Kmeans會(huì)直接調(diào)sklearn，如果任務(wù)比較復(fù)雜，可以通過(guò)numpy進(jìn)行自定義，這里介紹使用Pytorch實(shí)現(xiàn)的方式，經(jīng)測(cè)試，通過(guò)Pytorch調(diào)用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 = device

    def fit(self, x):
        # 隨機(jī)選擇初始中心點(diǎn)，想更快的收斂速度可以借鑒sklearn中的kmeans++初始化方法
        init_row = torch.randint(0, x.shape[0], (self.n_clusters,)).to(self.device)
        init_points = x[init_row]
        self.centers = init_points
        while True:
            # 聚類標(biāo)記
            self.nearest_center(x)
            # 更新中心點(diǎn)
            self.update_center(x)
            if self.verbose:
                print(self.variation, torch.argmin(self.dists, (0)))
            if torch.abs(self.variation) < 1e-3 and self.max_iter is None:
                break
            elif self.max_iter is not None and self.count == self.max_iter:
                break

            self.count += 1

        self.representative_sample()

    def nearest_center(self, x):
        labels = torch.empty((x.shape[0],)).long().to(self.device)
        dists = torch.empty((0, self.n_clusters)).to(self.device)
        for i, sample in enumerate(x):
            dist = torch.sum(torch.mul(sample - self.centers, sample - self.centers), (1))
            labels[i] = torch.argmin(dist)
            dists = torch.cat([dists, dist.unsqueeze(0)], (0))
        self.labels = labels
        if self.started:
            self.variation = torch.sum(self.dists - dists)
        self.dists = dists
        self.started = True

    def update_center(self, x):
        centers = torch.empty((0, x.shape[1])).to(self.device)
        for i in range(self.n_clusters):
            mask = self.labels == i
            cluster_samples = x[mask]
            centers = torch.cat([centers, torch.mean(cluster_samples, (0)).unsqueeze(0)], (0))
        self.centers = centers

    def representative_sample(self):
        # 查找距離中心點(diǎn)最近的樣本，作為聚類的代表樣本，更加直觀
        self.representative_samples = torch.argmin(self.dists, (0))


def time_clock(matrix,device):
    a = time.time()
    k = KMEANS(max_iter=10,verbose=False,device=device)
    k.fit(matrix)
    b = time.time()
    return (b-a)/k.count

def choose_device(cuda=False):
    if cuda:
        device = torch.device("cuda:0")
    else:
        device = torch.device("cpu")
    return device

if __name__ == "__main__":
    import matplotlib.pyplot as plt

    plt.figure()

    device = choose_device(False)

    cpu_speeds = []
    for i in tqdm([20,100,500,2000,8000,20000]):
        matrix = torch.rand((10000,i)).to(device)
        speed = time_clock(matrix,device)
        cpu_speeds.append(speed)
    l1, = plt.plot([20,100,500,2000,8000,20000],cpu_speeds,color = 'r',label = 'CPU')

    device = choose_device(True)

    gpu_speeds = []
    for i in tqdm([20, 100, 500, 2000, 8000, 20000]):
        matrix = torch.rand((10000, i)).to(device)
        speed = time_clock(matrix,device)
        gpu_speeds.append(speed)
    l2, = plt.plot([20, 100, 500, 2000, 8000, 20000], gpu_speeds, color='g',label = "GPU")



    plt.xlabel("num_features")
    plt.ylabel("speed(s/iter)")
    plt.title("Speed with cuda")
    plt.legend(handles = [l1,l2],labels = ['CPU','GPU'],loc='best')
    plt.savefig("../result/speed.jpg")

cpu和gpu運(yùn)行的結(jié)果對(duì)比如下：

speed.jpg

可以看到，在特征數(shù)<3000的情況下，cpu運(yùn)行速度更快，但是特征數(shù)量超過(guò)3000之后，gpu的優(yōu)勢(shì)越來(lái)越明顯。

因?yàn)閜ytorch的矩陣運(yùn)算接口基本是照著numpy寫的，所以numpy的實(shí)現(xiàn)方式大概只需要將代碼中的torch替換成numpy就可以了。