python 多进程编程 之 加速kmeans算法

kmeans 算法代码来自@dz4543的https://blog.csdn.net/dz4543/article/details/80190177这篇文章，它重点介绍了kmeans的算法，和一份代码实现。实现结果是对一个图片像素进行聚类，生成一个新的有 k类颜色组成的图片。很炫酷的视觉效果。

 

我在自己 4核 电脑上跑了一下，耗时比较多做2分类大概需要60S，分成7类的话大概需要十多分钟了，所以想着用多进程来加速一下算法比较耗时的部分。于是有了这份多进程通讯的优化代码。

首先原博客代码是这样，主要算法部分如下：

    while change == True:
        change = False  # 重置
        start = time.time()

        ############### 第一部分 对 data 中的数据进行标记，标记结果记录在subCenter中
        for i in range(m):
            # 设置样本与聚类中心之间的最小的距离，初始值为正无穷
            minDist = np.inf
            minIndex = 0  # 所属的类别
            for j in range(k):
                # 计算i和每个聚类中心之间的距离
                dist = distance(data[i,], centroids[j,])
                if dist < minDist:
                    minDist = dist
                    minIndex = j
            # 判断是否需要改变
            if subCenter[i, 0] != minIndex:  # 需要改变
                change = True
                subCenter[i,] = np.mat([minIndex, minDist])
        cost = time.time()-start
        print("单进程对数据进行标记 耗时{}秒，处理数据长度 {}".format(cost,m))
        
        ################# 第二部分 重新计算聚类中心
        for j in range(k):
            sum_all = np.mat(np.zeros((1, n)))
            r = 0  # 每个类别中的样本的个数
            for i in range(m):
                if subCenter[i, 0] == j:  # 计算第j个类别
                    sum_all += data[i,]
                    r += 1
            for z in range(n):
                try:
                    # 所以centroids是一个2维矩阵，每行为对应类别的中心
                    centroids[j, z] = sum_all[0, z] / r
                except:
                    print(" r is zero")
    return subCenter, centroids

从功能来看主要分两部分，已经在备注中说明，我选择了第一部分代码进行优化。优化后的代码，kmeas主逻辑放在一个线程中执行。第一部分的计算交给一堆进程去计算。这两个部分通过 multiprocessing.Pipe 对象生成的连接的 send 和 recv方法来进行通信。

优化之后的代码开4个进程实测结果如下，可以看到CPU利用率确实变高了，并且打印出的第一部分耗时也从900多秒变成了400多秒，左边为单进程CPU使用情况和耗时，右边为多进程情况。在更多核的电脑上起更多进程，优化效果更明显。

以下是优化后的代码：

import threading

import numpy as np
import cv2
import time
from multiprocessing import Process, Pipe

def perf_time(func):
    def wrap(*args):
        start = time.time()
        result = func(*args)
        cost = time.time() - start
        print("{} 调用耗时 {} 毫秒".format(func.__name__, cost*1000))
        return result
    return wrap

# 数据导入
@perf_time
def load_data(file_path):
    '''导入数据
    input:  file_path(string):文件的存储位置
    output: data(mat):数据
    '''
    data = []
    # 读取图片，转码后为矩阵，大小一般为m×n×3
    img = cv2.imread(file_path)
    # 获得图片大小
    m, n, _ = img.shape

    # 把图片展开，铺平
    for i in range(m):
        for j in range(n):
            tmp = []
            data.append(img[i, j,])
    return np.mat(data), m, n


# 定义相似性的度量
def distance(vecA, vecB):
    '''计算vecA与vecB之间的欧式距离的平方
    input:  vecA(mat)A点坐标
            vecB(mat)B点坐标
    output: dist[0, 0](float)A点与B点距离的平方
    '''
    dist = (vecA - vecB) * (vecA - vecB).T
    return dist[0, 0]

@perf_time
def randCenter(data, k):
    '''随机初始化聚类中心
    input:  data(mat):训练数据
            k(int):类别个数
    output: centroids(mat):聚类中心
    '''
    # 属性的个数，也就是load_data里的n
    _, n = data.shape

    # 初始化k个聚类中心，设为0
    centroids = np.mat(np.zeros((k, n)))
    # 初始化聚类中心每一维的坐标
    for j in range(n):
        # 求出每种特征值最小值，在图像里就是RGB取值最小值
        minJ = np.min(data[:, j])
        # 特征值取值范围
        rangeJ = np.max(data[:, j]) - minJ
        # 在最大值和最小值之间随机初始化，公式为：c=min+rand(0,1)×(max-min)
        centroids[:, j] = minJ * np.mat(np.ones((k, 1))) \
                          + np.random.rand(k, 1) * rangeJ
    return centroids

@perf_time
def kmeans(data, k, centroids, conns, processCount, result):
    '''根据KMeans算法求解聚类中心
    input:  data(mat):训练数据
            k(int):类别个数
            centroids(mat):随机初始化的聚类中心
    output: centroids(mat):训练完成的聚类中心
            subCenter(mat):每一个样本所属的类别
    '''
    # m：样本的个数，n：特征的维度
    m, n = np.shape(data)
    # 初始化每一个样本所属的类别，subCenter用来记录类别与相似度
    subCenter = np.mat(np.zeros((m, 2)))
    # 判断是否需要重新计算聚类中心
    change = True

    while change == True:
        change = False  # 重置

        # 拆分数据发送给不同进程计算
        for i in range(processCount):
            print("send msg to process {}".format(i))
            conns[i][0].send((change, centroids, subCenter[i::processCount]))
        start = time.time()
        subCenter_i_s = []
        # 从进程中回收处理后的 subCenter数据
        for i in range(processCount):
            recv_change, subCenter_i = conns[i][0].recv()
            if recv_change == True:
                change = True
            subCenter_i_s.append(subCenter_i)
        cost =time.time() - start
        print("多线程对数据进行标记 耗时 {} 秒, 处理数据长度 {}".format(cost, m))
        # 组合 subCenter
        if processCount == 1:
            subCenter = subCenter_i_s[0]
        else:
            subCenter = np.concatenate((subCenter_i_s[0], subCenter_i_s[1]), 1)
            for i in range(2, processCount):
                tmp = np.concatenate((subCenter, subCenter_i_s[i]),1)
                subCenter = tmp
        subCenter = subCenter.reshape(m, 2)
        print("组合后 subCenter.shape={}".format(subCenter.shape))
        #对样本按照聚类中心进行分类
        # 重新计算聚类中心
        for j in range(k):
            sum_all = np.mat(np.zeros((1, n)))
            r = 0  # 每个类别中的样本的个数
            for i in range(m):
                if subCenter[i, 0] == j:  # 计算第j个类别
                    sum_all += data[i,]
                    r += 1
            try:
                # 所以centroids是一个2维矩阵，每行为对应类别的中心
                centroids[j, ] = sum_all[0, ] / r
            except:
                print(" r is zero")
    for i in range(processCount):
        conns[i][0].send((None,None,None))
    result.append(subCenter)
    result.append(centroids)
    # result = [subCenter, centroids]


def annotationProcess(t_id, k, data_len, data, conns, processCount ):
    print("annotationProcess in !!!")
    while True:
        change, centroids, subCenter = conns[t_id][1].recv()
        if change == None:
            print("process {} recv None, end now !!!".format(t_id))
            break
        print("process {} recv data, change = {}".format(t_id, change))
        start = time.time()
        sub_id = 0
        count = 0
        for i in range(t_id, data_len, processCount):
            # 设置样本与聚类中心之间的最小的距离，初始值为正无穷
            minDist = np.inf
            minIndex = 0  # 所属的类别
            for j in range(k):
                # 计算i和每个聚类中心之间的距离
                dist = distance(data[i,], centroids[j,])
                if dist < minDist:
                    minDist = dist
                    minIndex = j
            # 判断是否需要改变
            if subCenter[sub_id, 0] != minIndex:  # 需要改变
                change = True
                subCenter[sub_id,] = np.mat([minIndex, minDist])
            sub_id+=1
            count+=1
        conns[t_id][1].send((change, subCenter))
        cost = time.time()- start
        print("process {} 本次处理耗时 {} 秒, change = {}, 共处理数据长度 {} ".format(t_id, cost, change, count))

if __name__ == "__main__":
    processCount = 4
    k = 7
    file_path = 'city.jpg'
    # 导入数据
    data, m, n = load_data(file_path)

    # 随机初始化聚类中心
    centroids = randCenter(data, k, )

    conns = []
    for i in range(processCount):
        conn1, conn2 = Pipe()
        conns.append((conn1, conn2))


    result = []
    # 启动 1个kmeans 线程
    main_thread = threading.Thread(target=kmeans, args=(data, k, centroids, conns, processCount, result))
    main_thread.start()
    process_s = []
    # 启动 processCount 个进程进行计算
    for i in range(processCount):
        p = Process(target=annotationProcess, args=(i, k, m*n, data, conns, processCount))
        process_s.append(p)
    for p in process_s:
        p.start()
    for p in process_s:
        p.join()

    main_thread.join()
    # 聚类结果
    subCenter, centroids = result

    # 保存分割后的图片
    new_pic = np.zeros((m * n, 3))
    print(new_pic.shape)

    for i in range(m * n):
        for j in range(k):
            if subCenter[i, 0] == j:
                new_pic[i, :] = centroids[j, :]
    new = np.reshape(new_pic, (m, n, 3))
    cv2.imwrite('new_pic.jpg', new)