雅虎BigML團隊開源大數(shù)據(jù)分布式深度學習框架TensorFlowOnSpark
近幾年，深度學習發(fā)展的非常迅速。在雅虎，我們發(fā)現(xiàn)，為了從海量數(shù)據(jù)中獲得洞察力，需要部署分布式深度學習?，F(xiàn)有的深度學習框架常常要求為深度學習單獨設定集群，迫使我們要為一個機器學習流程(見下圖 1)創(chuàng)建多個程序。
設定獨立的集群則需要我們轉移大數(shù)據(jù)集，帶來了不必要的系統(tǒng)復雜性和端到端的學習延遲。
去年我們通過開發(fā)和公開 CaffeOnSpark 解決了 scaleout 的問題，我們開源的框架支持在相同的 Spark 和 Hadoop 集群上進行分布式深度學習和大數(shù)據(jù)處理。我們在雅虎內部使用 CaffeOnSpark 改善了我們的 NSFW 圖像檢測，自動從實況錄像中識別電競比賽片段等等。在社區(qū)大量有價值的反饋和貢獻下，CaffeOnSpark 已經(jīng)得到了更新，現(xiàn)在可以支持 LSTM，有了一個新的數(shù)據(jù)層，可以訓練與測試交錯，有了一個 Python API，和 Docker container 的部署。這些都提升了我們的用戶體驗。但是那些使用 TensorFlow 框架的人怎么辦?于是我們效仿了之前的做法，開發(fā)了 TensorFlowOnSpark。
TensorFlow 公開后，谷歌于 2016 年 4 月就開放了一個帶有分布式學習功能的增強版 TensorFlow。2016 年 10 月，TensorFlow 開始支持 HDFS。然而在谷歌云之外，用戶仍然需要一個 TensorFlow 應用的專用集群。TensorFlow 程序無法在現(xiàn)有的大數(shù)據(jù)集群上部署，這樣一來，那些想大規(guī)模使用這個技術的人就需要花更多的成本和時間。
為了打破這個限制，幾個社區(qū)項目將 TensorFlow 連接到 Spark 集群上。SparkNet 讓 Spark 執(zhí)行器獲得了可以運行 TensorFlow 網(wǎng)絡的能力。DataBricks 提出 tensorframe，用來使用 TensorFlow 程序操縱 Apache Spark 的數(shù)據(jù)幀。雖然這些方法都朝著正確的方向邁出了一步，但是我們檢查他們的代碼后發(fā)現(xiàn)，我們無法讓多個 TensorFlow 過程直接相互溝通，我們也無法實現(xiàn)異步分布式學習，并且我們需要在遷移現(xiàn)有的 tensorflow 程序上花大功夫。
TensorFlowOnSpark
我們的新框架，TensorFlowOnSpark(TFoS)，支持 TensorFlow 在 Spark 和 Hadoop 上的分布式運行。如上圖(圖 2)所示，TFoS 與 SparkSQL、MLlib 以及其他的 Spark 庫一起在一個項目或線程(pipeline)中運行。
TFoS 支持所有類型的 TensorFlow 程序，能實現(xiàn)同步和異步的訓練與推理。并且支持模型和數(shù)據(jù)的平行處理，以及 TensorFlow 工具(如 TensorBoard)在 Spark 群集上使用。
任何 TensorFlow 程序都能夠很容易通過修改實現(xiàn)在 TFoS 上運行的。通常情況下，只需要修改少于 10 行的 Python 代碼。很多在雅虎平臺上使用 TensorFlow 的開發(fā)者，已經(jīng)輕松將 TensorFlow 項目轉移到 TFoS 上執(zhí)行了。
TFoS 支持張量(tensor)在 TensorFlow 處理過程中(計算節(jié)點和參數(shù)服務節(jié)點)信息的直接溝通。過程到過程(Process-to-process)的直接溝通機制使 TFoS 項目很容易在增加的機器上進行擴展。如圖 3 所示，TFoS 不需要 Spark 驅動器(driver)參與到張量溝通中來，因此也就與具備類似于獨立 TensorFlow 群集的擴展能力。
TFoS 提供兩種不同模式來「吞入」用于訓練和推理的數(shù)據(jù) ：
1. TensorFlow QueueRunners：TFoS 利用 TensorFlow 的文件讀取(file readers)和 QueueRunners 來直接從 HDFS 文件中讀入數(shù)據(jù)。在數(shù)據(jù)獲取過程中不需要 Spark 參與。
2. Spark 供給：Spark RDD 數(shù)據(jù)將會被傳輸至每一個 Spark 執(zhí)行器里，Spark 執(zhí)行器會進一步將數(shù)據(jù)傳入 TensorFlow 圖(通過 feed_dict 參數(shù))。
圖 4 展示了 Inception 圖像分類網(wǎng)絡中同時進行的分布式訓練如何在 TFoS 中通過 QueueRunners 的一個簡單設置進行擴展：將每個計算節(jié)點設置為 1 個 GPU，一個讀入(reader)以及批處理數(shù)為 32。四個 TFoS 的任務同時進行以用于訓練 10 萬步。兩天多后，當這些任務完成時，top-5 精確度(accuracy)分別為 0.730, 0.814, 0.854，0.879。0.730 的精確度需要單計算節(jié)點運行 46 小時得到，雙計算節(jié)點需要 22.5 個小時，4 計算機點需要 13 小時，8 計算節(jié)點需要 7.5 個小時。在 Inception 模型訓練上，TFoS 幾乎能達到線性擴展。這是很鼓舞人心的，雖然 TFoS 在不同模型和超參數(shù)上的擴展能力不同。
用于分布式 TensorFlow 的 RDMA
在雅虎的 Hadoop 集群上，GPU 節(jié)點通過以太網(wǎng)和無線寬帶相互連接。無線寬帶提供了高速的連接，并支持在 RDMA 中直接訪問其他服務器的存儲。然而目前 TensorFlow 僅支持在以太網(wǎng)上使用 「gRPC」 的分布式學習。為了加速分布式學習，我們增強了 TensorFlowC++層，實現(xiàn)了無線寬帶上的 RDMA。
為了結合我們發(fā)布的 TFoS，我們在 default「gRPC」協(xié)議之外，引進了一個新的 TensorFlow 服務器協(xié)議。任何分布式 tensorflow 程序都能通過指定 protocol=「grpc_rdma」in tf.train.ServerDef()or tf.train.Server() 來使用我們的增強版的 TensorFlow。
有了這個新協(xié)議后，就需要一個 RDMA 匯集管理器(rendezvous manager)來確保張量直接寫入遠程服務器的內存。我們最大限度地減少張量緩沖的創(chuàng)建：張量緩沖器在開始時分配一次，然后在一個 TensorFlow 工作任務的所用訓練步驟中重復使用。從我們早期的大型模型實驗，比如 VGG-19 開始，我們的就已經(jīng)證明了，與現(xiàn)有的 gRPC 相比，我們的 RDMA 實現(xiàn)在訓練時間上帶來了顯著的提速。
由于 RDMA 支持對性能要求很高(見 TensorFlow issue#2916)，我們決定讓我們現(xiàn)有的實現(xiàn)版本作為一個預覽版向 TensorFlow 社區(qū)開放。在未來的幾周內，我們將會進一步優(yōu)化我們的 RDMA 實現(xiàn)，并分享一些基準結果細節(jié)。
簡單的 CLI 和 API
TFoS 程序是通過標準的 ApacheSpark 命令 spark-submit 運行的。如下所示，用戶可以在 CLI 中指定 Spark 執(zhí)行器的數(shù)目，每個執(zhí)行器所用的 GPU 數(shù)目以及參數(shù)服務節(jié)點數(shù)。用戶還可以表明愿意使用 TensorBoard(–tensorboard)還是 RDMA(–rdma)。
- spark-submit –master ${MASTER} \
- ${TFoS_HOME}/examples/slim/train_image_classifier.py \
- –model_name inception_v3 \
- –train_dir hdfs://default/slim_train \
- –dataset_dir hdfs://default/data/imagenet \
- –dataset_name imagenet \
- –dataset_split_name train \
- –cluster_size ${NUM_EXEC} \
- –num_gpus ${NUM_GPU} \
- –num_ps_tasks ${NUM_PS} \
- –sync_replicas \
- –replicas_to_aggregate ${NUM_WORKERS} \
- –tensorboard \
- –rdma
TFoS 提供高層次的 Python API(在 Python notebook 的范例中有顯示)：
- TFCluster.reserve()... 從 Spark 執(zhí)行器構建一個 TensorFlow 群集
- TFCluster.start()... 在執(zhí)行器上加載 TensorFlow 程序
- TFCluster.train() or TFCluster.inference() …將 RDD 數(shù)據(jù)傳入 TensorFlow 處理
- TFCluster.shutdown() …在執(zhí)行器中結束 TensorFlow 的運行
開源
- TensorFlowOnSpark 開源地址： github.com/yahoo/TensorFlowOnSpark
- RDMA 增強版開源地址： github.com/yahoo/tensorflow/tree/yahoo
- 提供多示例程序(包括 MNIST，Cifar10，Inception，and VGG)以說明 TensorFlow 程序到TensorFlowOnspar轉換過程，并且利用 RDMA。地址：https://github.com/yahoo/TensorFlowOnSpark/tree/master/examples
- 提供一張亞馬遜機器圖像用于在 AWS EC2 上應用 TensorFlowOnSpark。接著，與 CaffeOnSpark 一樣，我們會推進 TensorFlowOnSpark。地址：https://github.com/yahoo/TensorFlowOnSpark/wiki/GetStarted_EC2
【本文是51CTO專欄機構機器之心的原創(chuàng)文章，微信公眾號“機器之心( id: almosthuman2014)”】
