在本文中，Ibotta（美國版「返利網(wǎng)」）機(jī)器學(xué)習(xí)和數(shù)據(jù)科學(xué)經(jīng)理 Evan Harris 介紹了他們的開源項目 sk-dist。這是一個分配 scikit-learn 元估計器的 Spark 通用框架，它結(jié)合了 Spark 和 scikit-learn 中的元素，可以將 sklearn 的訓(xùn)練速度提升 100 多倍。

在 Ibotta，我們訓(xùn)練了許多機(jī)器學(xué)習(xí)模型。這些模型為我們的推薦系統(tǒng)、搜索引擎、定價優(yōu)化引擎、數(shù)據(jù)質(zhì)量等提供了支持，在與我們的移動 app 互動的同時為數(shù)百萬用戶做出預(yù)測。

雖然我們使用 Spark 進(jìn)行大量的數(shù)據(jù)處理，但我們首選的機(jī)器學(xué)習(xí)框架是 scikit-learn。隨著計算成本越來越低以及機(jī)器學(xué)習(xí)解決方案的上市時間越來越重要，我們已經(jīng)踏出了加速模型訓(xùn)練的一步。其中一個解決方案是將 Spark 和 scikit-learn 中的元素組合，變成我們自己的融合解決方案。

項目地址：https://github.com/Ibotta/sk-dist

何為 sk-dist

我們很高興推出我們的開源項目 sk-dist。該項目的目標(biāo)是提供一個分配 scikit-learn 元估計器的 Spark 通用框架。元估計器的應(yīng)用包括決策樹集合（隨機(jī)森林和 extra randomized trees）、超參數(shù)調(diào)優(yōu)（網(wǎng)格搜索和隨機(jī)搜索）和多類技術(shù)（一對多和一對一）。

我們的主要目的是填補(bǔ)傳統(tǒng)機(jī)器學(xué)習(xí)模型分布選擇空間的空白。在神經(jīng)網(wǎng)絡(luò)和深度學(xué)習(xí)的空間之外，我們發(fā)現(xiàn)訓(xùn)練模型的大部分計算時間并未花在單個數(shù)據(jù)集上的單個模型訓(xùn)練上，而是花在用網(wǎng)格搜索或集成等元估計器在數(shù)據(jù)集的多次迭代中訓(xùn)練模型的多次迭代上。

實(shí)例

以手寫數(shù)字?jǐn)?shù)據(jù)集為例。我們編碼了手寫數(shù)字的圖像以便于分類。我們可以利用一臺機(jī)器在有 1797 條記錄的數(shù)據(jù)集上快速訓(xùn)練一個支持向量機(jī)，只需不到一秒。但是，超參數(shù)調(diào)優(yōu)需要在訓(xùn)練數(shù)據(jù)的不同子集上進(jìn)行大量訓(xùn)練。

如下圖所示，我們已經(jīng)構(gòu)建了一個參數(shù)網(wǎng)格，總共需要 1050 個訓(xùn)練項。在一個擁有 100 多個核心的 Spark 集群上使用 sk-dist 僅需 3.4 秒。這項工作的總?cè)蝿?wù)時間是 7.2 分鐘，這意味著在一臺沒有并行化的機(jī)器上訓(xùn)練需要這么長的時間。

import timefrom sklearn import datasets, svm 
from skdist.distribute.search import DistGridSearchCV 
from pyspark.sql import SparkSession # instantiate spark session 
spark = (    
    SparkSession     
    .builder     
    .getOrCreate()     
    ) 
sc = spark.sparkContext  
 
# the digits dataset 
digits = datasets.load_digits() 
X = digits["data"] 
y = digits["target"] 
 
# create a classifier: a support vector classifier 
classifier = svm.SVC() 
param_grid = { 
    "C": [0.01, 0.01, 0.1, 1.0, 10.0, 20.0, 50.0],  
    "gamma": ["scale", "auto", 0.001, 0.01, 0.1],  
    "kernel": ["rbf", "poly", "sigmoid"] 
    } 
scoring = "f1_weighted" 
cv = 10 
 
# hyperparameter optimization 
start = time.time() 
model = DistGridSearchCV(     
    classifier, param_grid,      
    sc=sc, cv=cv, scoring=scoring, 
    verbose=True     
    ) 
model.fit(X,y) 
print("Train time: {0}".format(time.time() - start)) 
print("Best score: {0}".format(model.best_score_)) 
 
 
------------------------------ 
Spark context found; running with spark 
Fitting 10 folds for each of 105 candidates, totalling 1050 fits 
Train time: 3.380601406097412 
Best score: 0.981450024203508 

該示例說明了一個常見情況，其中將數(shù)據(jù)擬合到內(nèi)存中并訓(xùn)練單個分類器并不重要，但超參數(shù)調(diào)整所需的擬合數(shù)量很快就會增加。以下是運(yùn)行網(wǎng)格搜索問題的內(nèi)在機(jī)制，如上例中的 sk-dist：

使用 sk-dist 進(jìn)行網(wǎng)格搜索

對于 Ibotta 傳統(tǒng)機(jī)器學(xué)習(xí)的實(shí)際應(yīng)用，我們經(jīng)常發(fā)現(xiàn)自己處于類似情況：中小型數(shù)據(jù)（100k 到 1M 記錄），其中包括多次迭代的簡單分類器，適合于超參數(shù)調(diào)優(yōu)、集合和多類解決方案。

現(xiàn)有解決方案

對于傳統(tǒng)機(jī)器學(xué)習(xí)元估計訓(xùn)練，現(xiàn)有解決方案是分布式的。第一個是最簡單的：scikit-learn 使用 joblib 內(nèi)置元估計器的并行化。這與 sk-dist 非常相似，除了一個主要限制因素：性能受限。即使對于具有數(shù)百個內(nèi)核的理論單臺機(jī)器，Spark 仍然具有如執(zhí)行器的內(nèi)存調(diào)優(yōu)規(guī)范、容錯等優(yōu)點(diǎn)，以及成本控制選項，例如為工作節(jié)點(diǎn)使用 Spot 實(shí)例。

另一個現(xiàn)有的解決方案是 Spark ML。這是 Spark 的本機(jī)機(jī)器學(xué)習(xí)庫，支持許多與 scikit-learn 相同的算法，用于分類和回歸問題。它還具有樹集合和網(wǎng)格搜索等元估計器，以及對多類問題的支持。雖然這聽起來可能是分配 scikit-learn 模式機(jī)器學(xué)習(xí)工作負(fù)載的優(yōu)秀解決方案，但它的分布式訓(xùn)練并不能解決我們感興趣的并行性問題。

分布在不同維度

如上所示，Spark ML 將針對分布在多個執(zhí)行器上的數(shù)據(jù)訓(xùn)練單個模型。當(dāng)數(shù)據(jù)很大且無法將內(nèi)存放在一臺機(jī)器上時，這種方法非常有效。但是，當(dāng)數(shù)據(jù)很小時，它在單臺計算機(jī)上的表現(xiàn)可能還不如 scikit-learn。此外，當(dāng)訓(xùn)練隨機(jī)森林時，Spark ML 按順序訓(xùn)練每個決策樹。無論分配給任務(wù)的資源如何，此任務(wù)的掛起時間都將與決策樹的數(shù)量成線性比例。

對于網(wǎng)格搜索，Spark ML 確實(shí)實(shí)現(xiàn)了并行性參數(shù)，將并行訓(xùn)練單個模型。但是，每個單獨(dú)的模型仍在對分布在執(zhí)行器中的數(shù)據(jù)進(jìn)行訓(xùn)練。如果按照模型的維度而非數(shù)據(jù)進(jìn)行分布，那么任務(wù)的總并行度可能是它的一小部分。

最終，我們希望將我們的訓(xùn)練分布在與 Spark ML 不同的維度上。使用小型或中型數(shù)據(jù)時，將數(shù)據(jù)擬合到內(nèi)存中不是問題。對于隨機(jī)森林的例子，我們希望將訓(xùn)練數(shù)據(jù)完整地廣播給每個執(zhí)行器，在每個執(zhí)行器上擬合一個獨(dú)立的決策樹，并將那些擬合的決策樹返回驅(qū)動程序以構(gòu)建隨機(jī)森林。沿著這個維度分布比串行分布數(shù)據(jù)和訓(xùn)練決策樹快幾個數(shù)量級。這種行為與網(wǎng)格搜索和多類等其他元估計器技術(shù)類似。

特征

鑒于這些現(xiàn)有解決方案在我們的問題空間中的局限性，我們決定在內(nèi)部開發(fā) sk-dist。最重要的是我們要「分配模型，而非數(shù)據(jù)」。

sk-dist 的重點(diǎn)是關(guān)注元估計器的分布式訓(xùn)練，還包括使用 Spark 進(jìn)行 scikit-learn 模型分布式預(yù)測的模塊、用于無 Spark 的幾個預(yù)處理/后處理的 scikit-learn 轉(zhuǎn)換器以及用于有/無 Spark 的靈活特征編碼器。

分布式訓(xùn)練：使用 Spark 分配元估計器訓(xùn)練。支持以下算法：超參數(shù)調(diào)優(yōu)（網(wǎng)格搜索和隨機(jī)搜索）、決策樹集合（隨機(jī)森林、額外隨機(jī)樹和隨機(jī)樹嵌入）以及多類技術(shù)（一對多和一對一）。

分布式預(yù)測：使用 Spark DataFrames 分布擬合 scikit-learn 估算器的預(yù)測方法?？梢酝ㄟ^便攜式 scikit-learn 估計器實(shí)現(xiàn)大規(guī)模分布式預(yù)測，這些估計器可以使用或不使用 Spark。

特征編碼：使用名為 Encoderizer 的靈活特征轉(zhuǎn)換器分布特征編碼。它可以使用或不使用 Spark 并行化。它將推斷數(shù)據(jù)類型和形狀，自動應(yīng)用默認(rèn)的特征轉(zhuǎn)換器作為標(biāo)準(zhǔn)特征編碼技術(shù)的最佳預(yù)測實(shí)現(xiàn)。它還可以作為完全可定制的特征聯(lián)合編碼器使用，同時具有與 Spark 分布式轉(zhuǎn)換器配合的附加優(yōu)勢。

用例

以下是判斷 sk-dist 是否適合你的機(jī)器學(xué)習(xí)問題空間的一些指導(dǎo)原則：

傳統(tǒng)機(jī)器學(xué)習(xí) ：廣義線性模型、隨機(jī)梯度下降、最近鄰算法、決策樹和樸素貝葉斯適用于 sk-dist。這些都可在 scikit-learn 中實(shí)現(xiàn)，可以使用 sk-dist 元估計器直接實(shí)現(xiàn)。

中小型數(shù)據(jù) ：大數(shù)據(jù)不適用于 sk-dist。請記住，訓(xùn)練分布的維度是沿著模型變化，而不是數(shù)據(jù)。數(shù)據(jù)不僅需要適合每個執(zhí)行器的內(nèi)存，還要小到可以廣播。根據(jù) Spark 配置，最大廣播大小可能會受到限制。

Spark 定位與訪問：sk-dist 的核心功能需要運(yùn)行 Spark。對于個人或小型數(shù)據(jù)科學(xué)團(tuán)隊而言，這并不總是可行的。此外，為了利用 sk-dist 獲得最大成本效益，需要進(jìn)行一些 Spark 調(diào)整和配置，這需要對 Spark 基礎(chǔ)知識進(jìn)行一些訓(xùn)練。

這里一個重要的注意事項是，雖然神經(jīng)網(wǎng)絡(luò)和深度學(xué)習(xí)在技術(shù)上可以與 sk-dist 一起使用，但這些技術(shù)需要大量的訓(xùn)練數(shù)據(jù)，有時需要專門的基礎(chǔ)設(shè)施才能有效。深度學(xué)習(xí)不是 sk-dist 的預(yù)期用例，因為它違反了上面的 (1) 和 (2)。在 Ibotta，我們一直在使用 Amazon SageMaker 這些技術(shù)，我們發(fā)現(xiàn)這些技術(shù)對這些工作負(fù)載的計算比使用 Spark 更有效。