基于TensorFlow的K-means算法的研究

李昱鋒，李建宏，文永明

(1.華北計(jì)算機(jī)系統(tǒng)工程研究所，北京 100083；2.中國人民解放軍第5718工廠，廣西 桂林 541003)

0 引言

聚類[1]分析是一個(gè)數(shù)據(jù)對象劃分成子集的過程，每個(gè)子集是一個(gè)簇，使得簇中的對象彼此相似，但與其他簇中的對象不相似。聚類分析已經(jīng)廣泛地用于許多應(yīng)用領(lǐng)域，包括商務(wù)智能、圖像模式識別、Web搜索、生物學(xué)和安全。在商務(wù)智能應(yīng)用中，可以將大量客戶分組，其中組內(nèi)的客戶具有非常類似的特征，這有利于開發(fā)加強(qiáng)客戶關(guān)系管理的商務(wù)策略。K-means算法是聚類分析中的一個(gè)經(jīng)典算法，是基于距離的聚類分析。在當(dāng)今信息時(shí)代，數(shù)據(jù)爆炸式增長，數(shù)據(jù)處理的計(jì)算量也隨之增大，為了縮短運(yùn)算時(shí)長采用GPU進(jìn)行計(jì)算，通過計(jì)算并行化從而提高算法性能。CUDA是NVIDIA推出的基于GPU的通用并行計(jì)算平臺和編程模型。TensorFlow支持GPU加速，通過使用CUDA和驅(qū)動實(shí)現(xiàn)GPU計(jì)算。

1 基于TensorFlow的K-means算法

1.1 CUDA平臺

CPU由專為順序串行處理而優(yōu)化的幾個(gè)核心組成，而GPU則擁有由數(shù)以千計(jì)的更小、更高效的核心(專為同時(shí)處理多重任務(wù)而設(shè)計(jì))組成的大規(guī)模并行計(jì)算架構(gòu)。這使得CPU更擅長處理具有復(fù)雜計(jì)算步驟和復(fù)雜數(shù)據(jù)依賴的計(jì)算任務(wù)，GPU更擅長處理并行計(jì)算[2]，CPU優(yōu)化復(fù)雜邏輯任務(wù)的運(yùn)行時(shí)間，GPU則優(yōu)化吞吐量。圖1為CPU和CPU的結(jié)構(gòu)圖。

圖1 CPU和GPU結(jié)構(gòu)圖

CUDA[3]是利用NVIDIA GPU的并行計(jì)算引擎以比CPU更高效的方式解決很多復(fù)雜的計(jì)算問題的通用并行計(jì)算平臺和編程模型。

CUDA是可擴(kuò)展的編程模型[4]，這種模型允許GPU架構(gòu)通過簡單擴(kuò)展多處理器和內(nèi)存分區(qū)的數(shù)量包含價(jià)格不同的設(shè)備，如圖2所示。

圖2 CUDA平臺擴(kuò)展圖

GPU是圍繞一系列流式多處理器(SM)構(gòu)建的。多線程程序被劃分為彼此獨(dú)立執(zhí)行的線程塊，因此具有更多多處理器的GPU比具有更少多處理器的GPU能在更短的時(shí)間內(nèi)執(zhí)行完程序。

1.2 CUDA編程模型

CUDA的C語言編程接口繼承了C語言中定義函數(shù)的方式來定義kernel，它不像C語言的函數(shù)只執(zhí)行一次，而是由N個(gè)不同CUDA線程并行執(zhí)行N次。一個(gè)kernel的定義由_global_聲明符號和CUDA線程數(shù)組成，線程數(shù)由特定的kernel調(diào)用通過使用<<<…>>>中的執(zhí)行參數(shù)配置。每一個(gè)執(zhí)行kernel的線程擁有唯一特定的thread ID，它在kernel內(nèi)部可通過訪問內(nèi)置的threadIdx變量得到。

threadIdx是一個(gè)三分量的向量，因此線程可以被一維、二維或者三維的線程索引所識別，形成一維、二維或者三維的線程塊。這提供了一種自然的方式調(diào)用向量、矩陣或者張量中域的元素進(jìn)行計(jì)算。線程的索引和線程ID對應(yīng)關(guān)系：對一維而言，它們相同；對大小為(Dx，Dy)的二維塊，線程索引為(x,y)的線程ID為(x+yDx)；對于大小為(Dx，Dy，Dz)的三維塊，線程索引為(x,y,z)的線程ID是(x+yDx+zDxDy)。塊中線程數(shù)目是有限制的，因?yàn)閴K中每個(gè)線程需要駐留在同一處理器核上并且共享有限的內(nèi)存資源，在目前的GPU上線程塊最多包含1 024個(gè)線程。然而，每個(gè)kernel能夠執(zhí)行多個(gè)同形狀的線程塊，所以總線程數(shù)目是塊數(shù)乘每塊的線程數(shù)。

線程塊需要獨(dú)立執(zhí)行：能夠以任何順序、并行或串行執(zhí)行它們。這種獨(dú)立性要求允許線程塊以任意順序在任意數(shù)目的核上進(jìn)行調(diào)度，使程序的代碼能夠具有在核數(shù)目上的擴(kuò)展性。同一個(gè)塊內(nèi)的線程通過共享內(nèi)存、數(shù)據(jù)和同步它們執(zhí)行協(xié)調(diào)內(nèi)存訪問實(shí)現(xiàn)合作。在kernel中同步的位置需要指定_syncthreads()，_syncthreads()作為一個(gè)塊內(nèi)所有線程允許繼續(xù)執(zhí)行前等待的阻塞點(diǎn)。為了協(xié)作的高效，共享內(nèi)存應(yīng)該是靠近每個(gè)處理器核的低延遲存儲器(類似L1 cache)，_syncthreads()應(yīng)該是輕量級的。

1.3 TensorFlow介紹

TensorFlow是一個(gè)用于研究和生產(chǎn)開發(fā)的開源機(jī)器學(xué)習(xí)庫。TensorFlow運(yùn)行時(shí)是一個(gè)跨平臺的庫，架構(gòu)如圖3所示。

圖3 TensorFlow架構(gòu)圖

Client：定義計(jì)算的數(shù)據(jù)流圖，使用session構(gòu)建圖的執(zhí)行。Distributed master：從圖中修剪由Session.run()中的參數(shù)定義特定的子圖，將子圖拆分成多個(gè)部分，它們運(yùn)行在不同的處理器和設(shè)備上；分發(fā)圖的各部分給work services，通過work services啟動圖各部分的執(zhí)行。Worker Services：使用合適可用的硬件(CPUs、GPUs等)的內(nèi)核實(shí)現(xiàn)調(diào)度圖運(yùn)算的執(zhí)行；發(fā)送和接收運(yùn)算結(jié)果給其它work services。Kernel implementation：執(zhí)行單個(gè)圖操作的計(jì)算。

TensorFlow中的op由kernel實(shí)現(xiàn)，GPU的內(nèi)核由兩部分組成：OpKernel和CUDA內(nèi)核以及啟動代碼。

TensorFlow這一框架定義和運(yùn)行涉及張量的計(jì)算。張量是矢量和矩陣向可能更高維的泛化。TensorFlow在內(nèi)部將張量表示為基本數(shù)據(jù)類型的n維數(shù)組。TensorFlow程序主要的操作和傳遞的對象是張量，一個(gè)張量對象表示部分定義為會最終產(chǎn)生值的計(jì)算。TensorFlow程序工作時(shí)首先要構(gòu)造一個(gè)圖的張量對象，每個(gè)張量的計(jì)算細(xì)節(jié)基于其他可用的張量和運(yùn)行該圖的部分以獲得所需的結(jié)果。張量具有兩個(gè)屬性：數(shù)據(jù)類型和形狀。張量中的每個(gè)元素都具有相同的數(shù)據(jù)類型，且該數(shù)據(jù)類型是已知的，形狀，張量的維數(shù)和每個(gè)維度的大小。

TensorFlow中的變量是程序操作共享和持久狀態(tài)的最好方式。通過tf.Variable類對變量進(jìn)行操作。tf.Variable相當(dāng)于一個(gè)張量的值，可以通過在其上運(yùn)行ops進(jìn)行改變。不同于張量對象，變量存在單個(gè)session.run調(diào)用的上下文之外。一個(gè)tf.Variable可以存儲一個(gè)持久的張量，特定的ops可以去讀取和修改張量的值。這些修改在多個(gè)tf.Session間是可見的。

TensorFlow使用數(shù)據(jù)流圖表示計(jì)算中獨(dú)立運(yùn)算間的依賴。在低級別的編程模型中首先定義數(shù)據(jù)流圖，然后創(chuàng)建TensorFlow中的session在一組本地和遠(yuǎn)程設(shè)備上運(yùn)行圖的各部分。數(shù)據(jù)流圖是一種用于并行計(jì)算的常見編程模型。在數(shù)據(jù)流圖中，節(jié)點(diǎn)代表著計(jì)算單元，邊代表著計(jì)算消耗或產(chǎn)生的數(shù)據(jù)。數(shù)據(jù)流帶來的優(yōu)勢：并行處理，通過使用明確的邊代表運(yùn)算間的依賴，這樣便于系統(tǒng)去識別能夠并行執(zhí)行的運(yùn)算；分布式執(zhí)行，通過明確的邊代表運(yùn)算間流動的值，TensorFlow能夠?qū)⒊绦騽澐值讲煌瑱C(jī)器的不同設(shè)備上，它在設(shè)備間添加必要的通信和協(xié)調(diào)；編譯生成高效代碼，TensorFlow的編譯器使用數(shù)據(jù)流圖生成還行更快的代碼，例如將相鄰的運(yùn)算融合在一起；可移植性，數(shù)據(jù)流圖描述的模型不依賴于語言。

tf.Graph包含兩類相關(guān)信息：圖結(jié)構(gòu)和圖集合。圖結(jié)構(gòu)信息主要包含圖的節(jié)點(diǎn)和邊，表示各個(gè)運(yùn)算如何組合在一起，但不指定它們該如何被使用。圖集合信息，TensorFlow提供一種用于存儲tf.Graph中的元數(shù)據(jù)集合的通用機(jī)制。

1.4 K-means算法

K-means算法[5-9]把簇的形心定義為簇內(nèi)點(diǎn)的均值。算法的處理流程如下。首先，在數(shù)據(jù)集D中隨機(jī)地選擇k個(gè)對象，每個(gè)對象代表一個(gè)簇的初始均值或中心。對剩下的每個(gè)對象，根據(jù)其與各個(gè)簇的中心的歐式距離，將它分配到最相似的簇。然后，k-means算法迭代地改善聚類結(jié)果。對于每個(gè)簇，它使用上次迭代分配到該簇的對象，計(jì)算新的均值。然后，使用更新后的均值作為新的簇中心，重新分配所有對象。迭代繼續(xù)，直到分配穩(wěn)定，即本輪形成的簇與前一輪形成的簇相同。

用于劃分的k-means算法，其中每個(gè)簇的中心都用簇中所有對象的均值來表示。

輸入：

k：簇的數(shù)目；

D：包含n個(gè)對象的數(shù)據(jù)集；

輸出：k個(gè)簇的集合；

方法：

(1)從D中任意選擇k個(gè)對象作為初始簇的中心；

(2)repeat;

(3)根據(jù)對象到簇中心的距離，將每個(gè)對象分配到最相似的簇；

(4)更新簇的中心，即重新計(jì)算每個(gè)簇中對象的均值；

(5)until不再發(fā)送變化。

1.5 基于TensorFlow的K-means算法優(yōu)化

K-means算法計(jì)算量主要體現(xiàn)在計(jì)算數(shù)據(jù)集中對象到簇中心的距離和更新簇中心兩步，通過提高這兩步的并行化進(jìn)而優(yōu)化算法[10-12]。

由于硬件結(jié)構(gòu)的不同，GPU比CPU更擅長并行計(jì)算，在GPU設(shè)備上執(zhí)行并行計(jì)算，算法性能會更好。英偉達(dá)廠商提供了CUDA計(jì)算平臺，它是一種通用并行計(jì)算框架，方便開發(fā)人員使用GPU解決復(fù)雜計(jì)算問題。英偉達(dá)為了深度學(xué)習(xí)框架提出了cuDNN，是用于深度神經(jīng)網(wǎng)絡(luò)的GPU加速原語庫。TensorFlow通過cuDNN和CUDA實(shí)現(xiàn)GPU計(jì)算，TensorFlow平臺提供的編程模型和編程接口比CUDA平臺更容易理解和使用。TensorFlow的編程接口表達(dá)能力強(qiáng)，使GPU計(jì)算更簡單且無需關(guān)注CUDA編程多線程并行的細(xì)節(jié)。TensorFlow描述的模型與語言無關(guān)，具有更好的移植性，架構(gòu)靈活，支持分布式。

在傳統(tǒng)的串行的K-means算法中，每個(gè)對象要計(jì)算k(簇?cái)?shù)目)次，才能得到各簇的中心距離。在TensorFlow中，通過將數(shù)據(jù)集轉(zhuǎn)化為張量，然后重構(gòu)張量用于計(jì)算距離，在新的張量中，數(shù)據(jù)集中的每個(gè)對象出現(xiàn)k次，相應(yīng)的各簇中心的重構(gòu)張量中，簇中心的每個(gè)對象出現(xiàn)n(數(shù)據(jù)集中記錄數(shù)目)次。求數(shù)據(jù)集中記錄到簇中心的距離，轉(zhuǎn)化為對兩個(gè)張量的計(jì)算，便于算法并行執(zhí)行，提高性能。

根據(jù)劃分的結(jié)果重新計(jì)算各簇的中心，簇的劃分結(jié)果由距離張量分組返回最小的索引，根據(jù)劃分結(jié)果計(jì)算簇各屬性的和，統(tǒng)計(jì)各簇中元素的數(shù)目，從而求得簇中對象的均值。

基于TensorFlow的K-means算法的主要步驟：

輸入：

k：簇的數(shù)目；

D：包含n個(gè)對象的數(shù)據(jù)集。

輸出：k個(gè)簇的集合。

方法：

(1)從D中任意選擇k個(gè)對象作為初始簇的中心；

(2)將數(shù)據(jù)集和簇中心轉(zhuǎn)為張量；

(3)調(diào)整數(shù)據(jù)集張量(記錄數(shù)，屬性數(shù))和簇中心張量(簇?cái)?shù)目，屬性數(shù))形狀，得到它們的重構(gòu)張量，形狀為(記錄數(shù)，簇?cái)?shù)目，屬性數(shù))；

(4)repeat；

(5)通過數(shù)據(jù)集和簇中心的重構(gòu)張量計(jì)算距離，分配對象到距離最小的簇；

(6)根據(jù)分配簇的結(jié)果張量和數(shù)據(jù)集張量，更新簇中心，重新計(jì)算簇中對象的均值；

(7)until不再發(fā)送變化。

算法的整體流程圖如圖4所示。

圖4 K-means優(yōu)化算法流程圖

2 實(shí)驗(yàn)與分析

本文實(shí)驗(yàn)所使用的數(shù)據(jù)集是美國人口普查數(shù)據(jù)集，該數(shù)據(jù)集包含大量人口普查數(shù)據(jù)，數(shù)據(jù)中的元素維度為69，每個(gè)對象擁有68種屬性，每個(gè)維度的數(shù)據(jù)類型是整型。

2.1 實(shí)驗(yàn)環(huán)境

本實(shí)驗(yàn)平臺為Intel(R) Core(TM) i7-8750H CPU，系統(tǒng)內(nèi)存為16 GB，GeForce GTX 1050 Ti，顯存4 GB，顯卡計(jì)算能力6.1，CUDA 9.0，cuDNN 9.0，TensorFlow 1.12.0。

2.2 實(shí)驗(yàn)結(jié)果

實(shí)驗(yàn)設(shè)置聚類中簇的數(shù)目k為8，由于簇中心的對象初始是隨機(jī)的，導(dǎo)致聚類收斂即算法的迭代次數(shù)不確定。為了避免迭代次數(shù)不同對實(shí)驗(yàn)產(chǎn)生影響，選取算法執(zhí)行迭代時(shí)間的平均單次時(shí)間為實(shí)驗(yàn)對象。實(shí)驗(yàn)分別在Python環(huán)境中使用CPU計(jì)算運(yùn)行普通的K-means算法，在TensorFlow環(huán)境中使用CPU計(jì)算運(yùn)行優(yōu)化的算法，在TensorFlow環(huán)境中使用GPU計(jì)算運(yùn)行優(yōu)化的算法。實(shí)驗(yàn)結(jié)果如表1所示，其中數(shù)據(jù)集記錄數(shù)目分別為：1萬條、5萬條、10萬條、50萬條。每項(xiàng)數(shù)據(jù)是通過運(yùn)行10次算法求均值得到的結(jié)果。

表1 算法運(yùn)行時(shí)間比較(ms)

由實(shí)驗(yàn)結(jié)果分析，GPU的架構(gòu)比CPU更適合處理并行任務(wù)，本實(shí)驗(yàn)環(huán)境中的CPU是多核的，因此具有一定并行計(jì)算能力，單次計(jì)算速度CPU也比GPU快。當(dāng)數(shù)據(jù)集不大的情況下，GPU計(jì)算速度并未有明顯優(yōu)勢，隨著數(shù)據(jù)量增大GPU的優(yōu)勢也開始明顯增大。使用相同的計(jì)算設(shè)備，TensorFlow編程模型更利于程序執(zhí)行的并行化。

3 結(jié)論

本文在經(jīng)典的K-means算法的基礎(chǔ)上，利用GPU擅長并行計(jì)算的特性，研究基于TensorFlow的K-means算法。該算法能有效地處理大規(guī)模數(shù)據(jù)集的聚類問題，利用張量運(yùn)算的思想加速了距離和簇中對象均值的計(jì)算，進(jìn)一步提高算法的并行化。TensorFlow簡化基于GPU編程，有良好的移植性，能夠在不同平臺上運(yùn)行和使用不同的計(jì)算設(shè)備。TensorFlow架構(gòu)支持分布式，分割描述算法的計(jì)算圖得到的子圖能夠在不同設(shè)備上執(zhí)行。通過實(shí)驗(yàn)證明該方法可以充分利用GPU的特性，有效加速K-means算法，縮短算法的運(yùn)行時(shí)間。