如何打造支持千億維特徵的機器學習基礎架構平台？

「範式大學系列課程」第 15 篇文章：如何打造支持千億維特徵的機器學習基礎架構

本文是大數據雜談 7 月 21 日社群公開課分享整理，也是第四範式主題月的第三堂公開課內容。

大家好，我是第四範式的陳迪豪，目前負責先知機器學習平台的架構與實現。

今天很高興和大家分享《打造機器學習的基礎架構平台》的話題，主要會介紹機器學習底層原理和工程實現方面的內容，也歡迎大家會後多多交流。

基礎架構（Infrastructure）相比於大數據、雲計算、深度學習，並不是一個很火的概念，甚至很多程序員就業開始就在用 MySQL、Django、Spring、Hadoop 來開發業務邏輯，而沒有真正參與過基礎架構項目的開發。在機器學習領域也是類似的，藉助開源的 Caffe、TensorFlow 或者 AWS、Google CloudML 就可以實現諸多業務應用，但框架或平台可能因行業的發展而流行或者衰退，而追求高可用、高性能、靈活易用的基礎架構卻幾乎是永恆不變的。

Google 的王詠剛老師在《為什麼 AI 工程師要懂一點架構》提到，研究院並不能只懂演算法，演算法實現不等於問題解決，問題解決不等於現場問題解決，架構知識是工程師進行高效團隊協作的共同語言。Google 依靠強大的基礎架構能力讓 AI 研究領先於業界，工業界的發展也讓深度學習、Auto Machine Learning 成為可能，未來將有更多人關注底層的架構與設計。

因此，今天的主題就是介紹機器學習的基礎架構，包括以下的幾個方面：

1. 基礎架構的分層設計；
2. 機器學習的數值計算；
3. TensorFlow 的重新實現；
4. 分散式機器學習平台的設計。

第一部分 基礎架構的分層設計

大家想像一下，如果我們在 AWS 上使用編寫一個 TensorFlow 應用，究竟經過了多少層應用抽象？首先，物理伺服器和網路寬頻就不必說了，通過 TCP/IP 等協議的抽象，我們直接在 AWS 虛擬機上操作就和本地操作沒有區別。其次，操作系統和編程語言的抽象，讓我們可以不感知底層內存物理地址和讀寫磁碟的 System call，而只需要遵循 Python 規範編寫代碼即可。然後，我們使用了 TensorFlow 計算庫，實際上我們只需調用最上層的 Python API，底層是經過了 Protobuf 序列化和 swig 進行跨語言調研，然後通過 gRPC 或者 RDMA 進行通信，而最底層這是調用 Eigen 或者 CUDA 庫進行矩陣運算。

因此，為了實現軟體間的解耦和抽象，系統架構常常採用分層架構，通過分層來屏蔽底層實現細節，而每一個底層都相當於上層應用的基礎架構。

那麼我們如何在一個分層的世界中夾縫生存？

有人可能認為，既然有人實現了操作系統和編程語言，那麼我們還需要關注底層的實現細節嗎？這個問題沒有標準答案，不同的人在不同的時期會有不同的感受，下面我舉兩個例子。

在《為了 1% 情形，犧牲 99% 情形下的性能：蝸牛般的 Python 深拷貝》這篇文章中，作者介紹了 Python 標準庫中 copy.deep_copy() 的實現，1% 的情況是指在深拷貝時對象內部有可能存在引用自身的對象，因此需要在拷貝時記錄所有拷貝過的對象信息，而 99% 的場景下對象並不會直接應用自身，為了兼容 100% 的情況這個庫損失了 6 倍以上的性能。在深入了解 Python 源碼後，我們可以通過實現深拷貝演算法來解決上述性能問題，從而優化我們的業務邏輯。

另一個例子是阿里的楊軍老師在 Strata Data Conference 分享的《Pluto: 一款分散式異構深度學習框架》，裡面介紹到基於 TensorFlow 的 control_dependencies 來實現冷熱數據在 GPU 顯存上的置入置出，從而在用戶幾乎不感知的情況下極大降低了顯存的使用量。了解源碼的人可能發現了，TensorFlow 的 Dynamic computation graph，也就是 tensorflow/fold 項目，也是基於 control_dependencies 實現的，能在聲明式機器學習框架中實現動態計算圖也是不太容易。這兩種實現都不存在 TensorFlow 的官方文檔中，只有對源碼有足夠深入的了解才可能在功能和性能上有巨大的突破，因此如果你是企業內 TensorFlow 框架的基礎架構維護者，突破 TensorFlow 的 Python API 抽象層是非常有必要的。

大家在應用機器學習時，不知不覺已經使用了很多基礎架構的抽象，其中最重要的莫過於機器學習演算法本身的實現，接下來我們將突破抽象，深入了解底層的實現原理。

第二部分 機器學習的數值計算

機器學習，本質上是一系列的數值計算，因此 TensorFlow 定位也不是一個深度學習庫，而是一個數值計算庫。當我們聽到了香農熵、貝葉斯、反向傳播這些概念時，並不需要擔心，這些都是數學，而且可以通過計算機編程實現的。

接觸過機器學習的都知道 LR，一般是指邏輯回歸（Logistic regression），也可以指線性回歸（Linear regression），而前者屬於分類演算法，後者屬於回歸演算法。兩種 LR 都有一些可以調優的超參數，例如訓練輪數（Epoch number）、學習率（Learning rate）、優化器（Optimizer）等，通過實現這個演算法可以幫忙我們理解其原理和調優技巧。

下面是一個最簡單的線性回歸 Python 實現，模型是簡單的 y = w * x + b。

從這個例子大家可以看到，實現一個機器學習演算法並不依賴於 Scikit-learn 或者 TensorFlow 等類庫，本質上都是數值運算，不同語言實現會有性能差異而已。細心的朋友可能發現，為什麼這裡 w 的梯度（Gradient）是 -2 * x * (y – x * x –b)，而 b 的梯度這是 -2 * (y – w * x - b)，如何保證經過計算後 Loss 下降而準確率上升？這就是數學上保證了，我們定義了 Loss 函數（Mean square error）為 y – w * x - b 的平方，也就是說預測值越接近 y 的話 Loss 越小，目標變成求 Loss 函數在 w 和 b 的任意取值下的最小值，因此對 w 和 b 求偏導後就得到上面兩條公式。

如果感興趣，不妨看一下線性回歸下 MSE 求偏導的數學公式證明。

邏輯回歸與線性回歸類似，當由於是分類問題，因此需要對 w * x + b 的預測結果進行歸一化（Normalization），一般使用 Sigmoid 方法，在 Python 中可以通過 1.0 / (1 + numpy.exp(-x)) 這種方式實現。由於預測值不同，Loss 函數的定義也不同，求偏導得到的數值計算公式也不同，感興趣也可以看看我的公式推導。

大家可以看到最終求得的偏導是非常簡單的，用任何編程語言都可以輕易實現。但我們自己的實現未必是最高效的，為什麼不直接用 Scikit-learn、MXNet 這些開源庫已經實現好的演算法呢？

我們對這個演算法的理解，其實是在工程上使用它的一個很重要的基礎。例如在真實的業務場景下，一個樣本的特徵可能有百億甚至千億維，而通過前面的演算法我們了解到，LR 模型的大小和樣本特徵的維度是相同的，也就是說一個接受百億維特徵的模型，本身參數就有百億個，如果使用標準的雙精度浮點數保存模型參數，那麼百億維的模型參數部分至少要超過 40G，那麼千億維的特徵更是單機所無法載入的。

因此，雖然 Scikit-learn 通過 native 介面實現了高性能的 LR 演算法，但只能滿足在單機上訓練，而 MXNet 由於原生沒有支持 SpareTensor，對於超高維度的稀疏數據訓練效率是非常低的，TensorFlow 本身支持 SpareTensor 也支持模型並行，可以支持百億維特徵的模型訓練，但沒有針對 LR 優化效率也不是很高。在這種場景下，第四範式基於 Parameter server 實現了支持模型並行和數據並行的超高維度、高性能機器學習庫，在此基礎上的大規模 LR、GBDT 等演算法訓練效率才能滿足工程上的需求。

機器學習還有很多有意思的演算法，例如決策樹、SVM、神經網路、樸素貝葉斯等等，只需要部分數學理論基礎就可以輕易在工程上實現，由於篇幅關係這裡就不在贅述了。前面我們介紹的其實是機器學習中的命令式（Imperative）編程介面，我們把求偏導的公式提前推導出來，然後像其他編程腳本一樣根據代碼那順序執行，而我們知道 TensorFlow 提供的是一種聲明式（Declarative）的編程介面，通過描述計算圖的方式來延後和優化執行過程，接下來我們就介紹這方面的內容。

第三部分 TensorFlow 的重新實現

首先大家可能有疑問，我們需要需要重新實現 TensorFlow？TensorFlow 靈活的編程介面、基於 Eigen 和 CUDA 的高性能計算、支持分散式和 Hadoop HDFS 集成，這些都是個人甚至企業很難完全追趕實現的，而且即使需要命令式編程介面我們也可以使用 MXNet，並沒有強需求需要一個新的 TensorFlow 框架。

事實上，我個人在學習 TensorFlow 過程中，通過實現一個 TensorFlow-like 的項目，不僅驚嘆與其源碼和介面的設計精巧，也加深了對聲明式編程、DAG 實現、自動求偏導、反向傳播等概念的理解。甚至在 Benchmark 測試中發現，純 Python 實現的項目在線性回歸模型訓練中比 TensorFlow 快 22 倍，當然這是在特定場景下壓測得到的結果，主要原因是 TensorFlow 中存在 Python 與 C++ 跨語言的切換開銷。

這個項目就是 MiniFlow，一個實現了鏈式法則、自動求導、支持命令式編程和聲明式編程的數值計算庫，並且兼容 TensorFlow Python API。感興趣可以在這個地址參與開發 https://github.com/tobegit3hub/miniflow ，下面是兩者 API 對比圖。

了解 TensorFlow 和 MXNet（或者 NNVM）源碼的朋友可能知道，兩者都抽象了 Op、Graph、Placeholer、Variable 等概念，通過 DAG 的方式描述模型的計算流圖，因此我們也需要實現類似的功能介面。

與前面的 LR 代碼不同，基於 Graph 的模型允許用戶自定義 Loss 函數，也就是用戶可以使用傳統的 Mean square error，也可以自定義一個任意的數學公式作為 Loss 函數，這要求框架本身能夠實現自動求導的功能，而不是我們根據 Loss 函數預先實現了導數的計算方式。

那麼用戶可以定義的最小操作，也就是 Op，需要平台實現基本的運算元，例如 ConstantOp、AddOp、MultipleOp 等，而且用戶實現自定義運算元時可以加入自動求導的流程中，並不影響框架本身的訓練流程。參考 TensorFlow 的 Python 源碼，下面我們定義了 Op 的基類，所有的 Op 都應該實現 forward() 和 grad() 以便於模型訓練時自動求導，而且通過重載 Python 操作符可以為開發者提供更便利的使用介面。

那麼對於常量（ConstantOp）和變數（VariableOp），他們的正向運算就是得到的是本身的值，而求導時常量的導數為 0，求偏導的變數導數為 1，其他變數也為 0，具體代碼如下。

其實更重要的是，我們需要實現加（AddOp）、減（MinusOp）、乘（MultipleOp）、除（DivideOp）、平方（PowerOp）等運算元的正向運算和反向運算邏輯，然後根據鏈式法則，任何複雜的數學公式求導都可以簡化成這些基本運算元的求導。

例如加法和減法，我們知道兩個數加法的導數等於導數的加法，因此根據此數學原理，我們可以很容易實現 AddOp，而 MinusOp 實現類似就不贅述了。

而乘法和除法相對複雜，顯然兩個數乘法的導數不等於導數的乘法，例如 x 和 x 的平方，先導數後相乘得到 2x，先相乘後導數得到 3 倍 x 的平方。因此這是需要使用乘數法則，基本公式是，而代碼實現如下。

除法和平方的求導方式也是類似的，因為數學上已經證明，所以只需要編碼實現基本的正向和反向運算即可。由於篇幅有限，這裡不再細緻介紹 MiniFlow 的源碼實現了，感興趣可以通過上面的 Github 鏈接找到完整的源碼實現，下面再提供使用相同 API 介面實現的模型性能測試結果，對於小批量數據處理、需要頻繁切換 Python/C++ 運行環境的場景下 MiniFlow 會有更好的性能表現。

前面介紹了機器學習演算法和深度學習類庫的實現，並不是所有人都有能力去重寫或者優化這部分基礎架構的，很多時候我們都只是這些演算法的使用者，但從另一個角度，我們就需要維護一個高可用的計算平台來做機器學習的訓練和預測，下面將從這方面介紹如何打造分散式機器學習平台。

第四部分 分散式機器學習平台的設計

隨著大數據和雲計算的發展，實現一個高可用、分散式的機器學習平台成為一個基本需求。無論是 Caffe、TensorFlow，還是我們自研的高性能機器學習庫，都只是解決數值計算、演算法實現以及模型訓練的問題，對於任務的隔離、調度、Failover 都需要上層平台實現。

那麼設計一個針對機器學習全流程的基礎架構平台，需要涵蓋哪些功能呢？

首先，必須實現資源隔離。在一個共享底層計算資源的集群中，用戶提交的訓練任務不應該受到其他任務的影響，儘可能保證 CPU、內存、GPU 等資源隔離。如果使用 Hadoop 或 Spark 集群，默認就會在任務進程上掛載 cgroups，保證 CPU 和內存的隔離，而隨著 Docker 等容器技術的成熟，我們也可以使用 Kubernetes、Mesos 等項目來啟動和管理用戶實現的模型訓練任務。

其次，實現資源調度和共享。隨著通用計算的 GPU 流行，目前支持 GPU 調度的編排工具也越來越多，而部分企業內還存在著 GPU 專卡專用的情況，無法實現資源的動態調度和共享，這必然導致計算資源的嚴重浪費。在設計機器學習平台時，需要儘可能考慮通用的集群共享場景，例如同時支持模型訓練、模型存儲以及模型服務等功能，可以對標的典例就是 Google Borg 系統。

然後，平台需要有靈活的兼容性。目前機器學習業務發展迅速，針對不同場景的機器學習框架也越來越多，靈活的平台架構可以兼容幾乎所有主流的應用框架，避免基礎架構因為業務的發展而頻繁變化。目前 Docker 是一種非常合適的容器格式規範，通過編寫 Dockerfile 就可以描述框架的運行環境和系統依賴，在此基礎上我們可以在平台上實現了 TensorFlow、MXNet、Theano、CNTK、Torch、Caffe、Keras、Scikit-learn、XGBoost、PaddlePaddle、Gym、Neon、Chainer、PyTorch、Deeplearning4j、Lasagne、Dsstne、H2O、GraphLab 以及 MiniFlow 等框架的集成。

最後，需要實現機器學習場景下的 API 服務。針對機器學習的模型開發、模型訓練和模型服務三個主要流程，我們可以定義提交訓練任務、創建開發環境、啟動模型服務、提交離線預測任務等 API，用熟悉的編程語言來實現 Web service 介面。要實現一個 Google-like 的雲深度學習平台，大家可以參考下面這三個步驟。

當然，要實現一個涵蓋數據引入、數據處理、特徵工程以及模型評估功能的機器學習平台，我們還需要集成 HDFS、Spark、Hive 等大數據處理工具，實現類似 Azkaban、Oozie 的工作流管理工具，在易用性、低門檻方面做更多的工作。

總結

最後總結一下，機器學習的基礎架構包含了機器學習演算法、機器學習類庫以及機器學習平台等多個層次的內容。根據業務的需求，我們可以選擇特定的領域進行深入研究和二次開發，利用輪子和根據需求改造輪子同樣重要。

在機器學習與人工智慧非常流行的今天，希望大家也可以重視底層基礎架構，演算法研究員可以 理解更多工程的設計與實現，而研發工程師可以了解更多的演算法原理與優化，在合適的基礎架構平台上讓機器學習發揮更大的效益，真正應用的實際場景中。

答疑環節

Q 1：老師您好! 我的問題是 基礎架構在具體落地方面 有什麼建議？比如在雲上部署和虛擬化容器技術的使用？

陳迪豪：基礎架構其實包含多層次的內容，如果在雲端部署，可以考慮使用 AWS 或者 Google CloudML 等基礎服務，也可以基於 Kubernetes、TensorFlow 等開源項目這內部搭建機器學習平台，參考前面圖片提到的三個步驟，只需要實現簡單的 API 服務和容器調度任務即可。

Q 2：老師好，接觸（機器學習）之前需要深入學習 spark 嗎？

陳迪豪：機器學習演算法本身並依賴 Spark，我們可以自己實現基本的演算法，也可以使用 Scikit-learn、TensorFlow 等開源庫。當然在大部分業務場景中，我們還是需要 Hadoop、Spark 等大數據框架進行數據處理、特徵抽取等功能，因此掌握一定的大數據處理能力也是很有價值的。

Q 3：我有個問題，怎樣做到線下模型效果評估的自動化？

陳迪豪：這是個好問題，在線下我們一般會對測試數據集計算 AUC 等離線指標，來預估模型的效果，自動化方面我們有一套自學習的流程，大家也可以使用 Crontab 或者 Azkaban 等任務管理工具，對於新的測試集進行模型評估，這是純工程的問題了可以結合已有的服務架構來實現。

Q 4：文中介紹的機器學習基礎架構平台的搭建思路，是否就是先知平台的架構思路？先知平台是 SaaS 式的服務平台，如果我是想搭建公司內部的機器學習架構平台，思路也是跟文中描述一樣還是說有什麼差異？

陳迪豪：先知平台的底層也是分散式、高可用的基礎架構，也是這個思路，不過在易用性、低門檻方面做了更多的工作。目前先知提供 SaaS 公有雲版本可以很方便得註冊使用，如果是企業內部搭建，先知平台也有企業版私有雲可以單獨部署，如果需要自己維護和搭建機器學習平台，可以考慮基於 Kubernetes 的容器調度集群，可以通過管理 CPU、GPU 等異構資源，通過製作 Docker 鏡像的方式來支持各種機器學習框架的使用，和前面提到的思路類似。

Q 5：對於一般的公司數據規模沒那麼大，是不是一個高配的機器 +4 個 GPU 卡，安裝一個 tensorflow 就可以了？

陳迪豪：對於規模比較小的公司，可以直接使用單機多卡的方式，包括 BAT 等大企業有的部門也是直接使用裸機的，但裸機帶來的問題是沒有資源管理和任務調度，使用者之間需要互相約定使用時間和使用的資源，在一定程度上造成資源浪費，如果規模大了建議使用統一的集群調度方式，可以保證任務間的資源隔離和正常執行。

Q 6：正如老師在群里上課提的，在工程上，在真實的業務場景下，一個樣本的特徵是百億甚至千億，在線下的時候，訓練好一個模型，部署在伺服器端。那麼，在線服務階段，也就是在線服務的預測階段的時候，由於發起預測請求的樣本特徵也比較大，那麼請問在實際工程中這些數據是怎麼存儲的？

陳迪豪：在一個百億甚至是千億維的模型上，我們一般通過分 shard 的方式來切分，模型本身可以通過多個 Parameter server 實例來保存，而客戶端發送預測請求時，由於特徵是非常稀疏的，實際有值的數據非常少，因此並不需要把大部分零值數據也請求到服務端，具體格式大家也可以參考下 TensorFlow 中 SparseTensor 的實現。

Q 7：請問先知目前平台的演算法是基於開源演算法封裝還是自己實現的呢？

陳迪豪：目前先知平台主要針對真實的應用場景，自主研發的大規模 LR、GBDT、HE-TreeNet、LFC、SVM、Feature-go 等演算法都集成到先知平台上，當然平台也支持 Spark MLlib、TensorFlow 等開源框架的演算法實現。

作者介紹

陳迪豪，第四範式先知平台架構師，曾在小米科技和 UnitedStack 擔任基礎架構研發工程師。活躍於 OpenStack、Kubernetes、TensorFlow 等開源社區，實現了 Cloud Machine Learning 雲深度學習平台，Github 賬號 https://github.com/tobegit3hub。

如果你對我們感興趣，請在這裡申請先知平台：https://prophet.4paradigm.com

AI FOR EVERYONE

推薦閱讀：