美國「返利網」開源sk-dist框架將sklearn訓練速度提升數倍

在本文中，Ibotta（美國版「返利網」）機器學習和資料科學經理 Evan Harris 介紹了他們的開源專案 sk-dist。

在本文中，Ibotta（美國版「返利網」）機器學習和資料科學經理 Evan Harris 介紹了他們的開源專案 sk-dist。這是一個分配 scikit-learn 元估計器的 Spark 通用框架，它結合了 Spark 和 scikit-learn 中的元素，可以將 sklearn 的訓練速度提升 100 多倍。

在 Ibotta，我們訓練了許多機器學習模型。這些模型為我們的推薦系統、搜尋引擎、定價最佳化引擎、資料質量等提供了支援，在與我們的移動 app 互動的同時為數百萬使用者做出預測。

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

我們很高興推出我們的開源專案 sk-dist。該專案的目標是提供一個分配 scikit-learn 元估計器的 Spark 通用框架。元估計器的應用包括決策樹集合（隨機森林和 extra randomized trees）、超引數調優（網格搜尋和隨機搜尋）和多類技術（一對多和一對一）。

我們的主要目的是填補傳統機器學習模型分佈選擇空間的空白。在神經網路和深度學習的空間之外，我們發現訓練模型的大部分計算時間並未花在單個資料集上的單個模型訓練上，而是花在用網格搜尋或整合等元估計器在資料集的多次迭代中訓練模型的多次迭代上。

以手寫數字資料集為例。我們編碼了手寫數字的影像以便於分類。我們可以利用一臺機器在有 1797 條記錄的資料集上快速訓練一個支援向量機，只需不到一秒。但是，超引數調優需要在訓練資料的不同子集上進行大量訓練。

如下圖所示，我們已經構建了一個引數網格，總共需要 1050 個訓練項。在一個擁有 100 多個核心的 Spark 叢集上使用 sk-dist 僅需 3.4 秒。這項工作的總任務時間是 7.2 分鐘，這意味著在一臺沒有並行化的機器上訓練需要這麼長的時間。

    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

該示例說明了一個常見情況，其中將資料擬合到記憶體中並訓練單個分類器並不重要，但超引數調整所需的擬合數量很快就會增加。以下是執行網格搜尋問題的內在機制，如上例中的 sk-dist：

對於 Ibotta 傳統機器學習的實際應用，我們經常發現自己處於類似情況：中小型資料（100k 到 1M 記錄），其中包括多次迭代的簡單分類器，適合於超引數調優、集合和多類解決方案。

現有解決方案

對於傳統機器學習元估計訓練，現有解決方案是分散式的。第一個是最簡單的：scikit-learn 使用 joblib 內建元估計器的並行化。這與 sk-dist 非常相似，除了一個主要限制因素：效能受限。即使對於具有數百個核心的理論單臺機器，Spark 仍然具有如執行器的記憶體調優規範、容錯等優點，以及成本控制選項，例如為工作節點使用 Spot 例項。

另一個現有的解決方案是 Spark ML。這是 Spark 的本機機器學習庫，支援許多與 scikit-learn 相同的演算法，用於分類和迴歸問題。它還具有樹集合和網格搜尋等元估計器，以及對多類問題的支援。雖然這聽起來可能是分配 scikit-learn 模式機器學習工作負載的優秀解決方案，但它的分散式訓練並不能解決我們感興趣的並行性問題。

分佈在不同維度
如上所示，Spark ML 將針對分佈在多個執行器上的資料訓練單個模型。當資料很大且無法將記憶體放在一臺機器上時，這種方法非常有效。但是，當資料很小時，它在單臺計算機上的表現可能還不如 scikit-learn。此外，當訓練隨機森林時，Spark ML 按順序訓練每個決策樹。無論分配給任務的資源如何，此任務的掛起時間都將與決策樹的數量成線性比例。

對於網格搜尋，Spark ML 確實實現了並行性引數，將並行訓練單個模型。但是，每個單獨的模型仍在對分佈在執行器中的資料進行訓練。如果按照模型的維度而非資料進行分佈，那麼任務的總並行度可能是它的一小部分。

最終，我們希望將我們的訓練分佈在與 Spark ML 不同的維度上。使用小型或中型資料時，將資料擬合到記憶體中不是問題。對於隨機森林的例子，我們希望將訓練資料完整地廣播給每個執行器，在每個執行器上擬合一個獨立的決策樹，並將那些擬合的決策樹返回驅動程式以構建隨機森林。沿著這個維度分佈比序列分佈資料和訓練決策樹快幾個數量級。這種行為與網格搜尋和多類等其他元估計器技術類似。

鑑於這些現有解決方案在我們的問題空間中的侷限性，我們決定在內部開發 sk-dist。最重要的是我們要「分配模型，而非資料」。

sk-dist 的重點是關注元估計器的分散式訓練，還包括使用 Spark 進行 scikit-learn 模型分散式預測的模組、用於無 Spark 的幾個預處理/後處理的 scikit-learn 轉換器以及用於有/無 Spark 的靈活特徵編碼器。

分散式訓練：使用 Spark 分配元估計器訓練。支援以下演算法：超引數調優（網格搜尋和隨機搜尋）、決策樹集合（隨機森林、額外隨機樹和隨機樹嵌入）以及多類技術（一對多和一對一）。

分散式預測：使用 Spark DataFrames 分佈擬合 scikit-learn 估算器的預測方法。可以透過行動式 scikit-learn 估計器實現大規模分散式預測，這些估計器可以使用或不使用 Spark。

特徵編碼：使用名為 Encoderizer 的靈活特徵轉換器分佈特徵編碼。它可以使用或不使用 Spark 並行化。它將推斷資料型別和形狀，自動應用預設的特徵轉換器作為標準特徵編碼技術的最佳預測實現。它還可以作為完全可定製的特徵聯合編碼器使用，同時具有與 Spark 分散式轉換器配合的附加優勢。

以下是判斷 sk-dist 是否適合你的機器學習問題空間的一些指導原則：

傳統機器學習 ：廣義線性模型、隨機梯度下降、最近鄰演算法、決策樹和樸素貝葉斯適用於 sk-dist。這些都可在 scikit-learn 中實現，可以使用 sk-dist 元估計器直接實現。

中小型資料 ：大資料不適用於 sk-dist。請記住，訓練分佈的維度是沿著模型變化，而不是資料。資料不僅需要適合每個執行器的記憶體，還要小到可以廣播。根據 Spark 配置，最大廣播大小可能會受到限制。

Spark 定位與訪問：sk-dist 的核心功能需要執行 Spark。對於個人或小型資料科學團隊而言，這並不總是可行的。此外，為了利用 sk-dist 獲得最大成本效益，需要進行一些 Spark 調整和配置，這需要對 Spark 基礎知識進行一些訓練。

這裡一個重要的注意事項是，雖然神經網路和深度學習在技術上可以與 sk-dist 一起使用，但這些技術需要大量的訓練資料，有時需要專門的基礎設施才能有效。深度學習不是 sk-dist 的預期用例，因為它違反了上面的 (1) 和 (2)。在 Ibotta，我們一直在使用 Amazon SageMaker 這些技術，我們發現這些技術對這些工作負載的計算比使用 Spark 更有效。

原文地址：