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

李昱鋒，李建宏，文永明

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

0 引言

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

1 基于TensorFlow的K-means算法

1.1 CUDA平臺

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

圖1 CPU和GPU結構圖

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

CUDA是可擴展的編程模型[4]，這種模型允許GPU架構通過簡單擴展多處理器和內存分區的數量包含價格不同的設備，如圖2所示。

圖2 CUDA平臺擴展圖

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

1.2 CUDA編程模型

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

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

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

1.3 TensorFlow介紹

TensorFlow是一個用于研究和生產開發的開源機器學習庫。TensorFlow運行時是一個跨平臺的庫，架構如圖3所示。

圖3 TensorFlow架構圖

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

TensorFlow中的op由kernel實現，GPU的內核由兩部分組成：OpKernel和CUDA內核以及啟動代碼。

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

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

TensorFlow使用數據流圖表示計算中獨立運算間的依賴。在低級別的編程模型中首先定義數據流圖，然后創建TensorFlow中的session在一組本地和遠程設備上運行圖的各部分。數據流圖是一種用于并行計算的常見編程模型。在數據流圖中，節點代表著計算單元，邊代表著計算消耗或產生的數據。數據流帶來的優勢：并行處理，通過使用明確的邊代表運算間的依賴，這樣便于系統去識別能夠并行執行的運算；分布式執行，通過明確的邊代表運算間流動的值，TensorFlow能夠將程序劃分到不同機器的不同設備上，它在設備間添加必要的通信和協調；編譯生成高效代碼，TensorFlow的編譯器使用數據流圖生成還行更快的代碼，例如將相鄰的運算融合在一起；可移植性，數據流圖描述的模型不依賴于語言。

tf.Graph包含兩類相關信息：圖結構和圖集合。圖結構信息主要包含圖的節點和邊，表示各個運算如何組合在一起，但不指定它們該如何被使用。圖集合信息，TensorFlow提供一種用于存儲tf.Graph中的元數據集合的通用機制。

1.4 K-means算法

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

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

輸入：

k：簇的數目；

D：包含n個對象的數據集；

輸出：k個簇的集合；

方法：

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

(2)repeat;

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

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

(5)until不再發送變化。

1.5 基于TensorFlow的K-means算法優化

K-means算法計算量主要體現在計算數據集中對象到簇中心的距離和更新簇中心兩步，通過提高這兩步的并行化進而優化算法[10-12]。

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

在傳統的串行的K-means算法中，每個對象要計算k(簇數目)次，才能得到各簇的中心距離。在TensorFlow中，通過將數據集轉化為張量，然后重構張量用于計算距離，在新的張量中，數據集中的每個對象出現k次，相應的各簇中心的重構張量中，簇中心的每個對象出現n(數據集中記錄數目)次。求數據集中記錄到簇中心的距離，轉化為對兩個張量的計算，便于算法并行執行，提高性能。

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

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

輸入：

k：簇的數目；

D：包含n個對象的數據集。

輸出：k個簇的集合。

方法：

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

(2)將數據集和簇中心轉為張量；

(3)調整數據集張量(記錄數，屬性數)和簇中心張量(簇數目，屬性數)形狀，得到它們的重構張量，形狀為(記錄數，簇數目，屬性數)；

(4)repeat；

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

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

(7)until不再發送變化。

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

圖4 K-means優化算法流程圖

2 實驗與分析

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

2.1 實驗環境

本實驗平臺為Intel(R) Core(TM) i7-8750H CPU，系統內存為16 GB，GeForce GTX 1050 Ti，顯存4 GB，顯卡計算能力6.1，CUDA 9.0，cuDNN 9.0，TensorFlow 1.12.0。

2.2 實驗結果

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

表1 算法運行時間比較(ms)

由實驗結果分析，GPU的架構比CPU更適合處理并行任務，本實驗環境中的CPU是多核的，因此具有一定并行計算能力，單次計算速度CPU也比GPU快。當數據集不大的情況下，GPU計算速度并未有明顯優勢，隨著數據量增大GPU的優勢也開始明顯增大。使用相同的計算設備，TensorFlow編程模型更利于程序執行的并行化。

3 結論

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