近幾年，深度學(xué)習(xí)發(fā)展的非常迅速。在雅虎，我們發(fā)現(xiàn)，為了從海量數(shù)據(jù)中獲得洞察力，需要部署分布式深度學(xué)習(xí)?，F(xiàn)有的深度學(xué)習(xí)框架常常要求為深度學(xué)習(xí)單獨(dú)設(shè)定集群，迫使我們要為一個(gè)機(jī)器學(xué)習(xí)流程(見下圖 1)創(chuàng)建多個(gè)程序。

設(shè)定獨(dú)立的集群則需要我們轉(zhuǎn)移大數(shù)據(jù)集，帶來了不必要的系統(tǒng)復(fù)雜性和端到端的學(xué)習(xí)延遲。

去年我們通過開發(fā)和公開 CaffeOnSpark 解決了 scaleout 的問題，我們開源的框架支持在相同的 Spark 和 Hadoop 集群上進(jìn)行分布式深度學(xué)習(xí)和大數(shù)據(jù)處理。我們?cè)谘呕?nèi)部使用 CaffeOnSpark 改善了我們的 NSFW 圖像檢測(cè)，自動(dòng)從實(shí)況錄像中識(shí)別電競(jìng)比賽片段等等。在社區(qū)大量有價(jià)值的反饋和貢獻(xiàn)下，CaffeOnSpark 已經(jīng)得到了更新，現(xiàn)在可以支持 LSTM，有了一個(gè)新的數(shù)據(jù)層，可以訓(xùn)練與測(cè)試交錯(cuò)，有了一個(gè) Python API，和 Docker container 的部署。這些都提升了我們的用戶體驗(yàn)。但是那些使用 TensorFlow 框架的人怎么辦?于是我們效仿了之前的做法，開發(fā)了 TensorFlowOnSpark。

TensorFlow 公開后，谷歌于 2016 年 4 月就開放了一個(gè)帶有分布式學(xué)習(xí)功能的增強(qiáng)版 TensorFlow。2016 年 10 月，TensorFlow 開始支持 HDFS。然而在谷歌云之外，用戶仍然需要一個(gè) TensorFlow 應(yīng)用的專用集群。TensorFlow 程序無法在現(xiàn)有的大數(shù)據(jù)集群上部署，這樣一來，那些想大規(guī)模使用這個(gè)技術(shù)的人就需要花更多的成本和時(shí)間。

為了打破這個(gè)限制，幾個(gè)社區(qū)項(xiàng)目將 TensorFlow 連接到 Spark 集群上。SparkNet 讓 Spark 執(zhí)行器獲得了可以運(yùn)行 TensorFlow 網(wǎng)絡(luò)的能力。DataBricks 提出 tensorframe，用來使用 TensorFlow 程序操縱 Apache Spark 的數(shù)據(jù)幀。雖然這些方法都朝著正確的方向邁出了一步，但是我們檢查他們的代碼后發(fā)現(xiàn)，我們無法讓多個(gè) TensorFlow 過程直接相互溝通，我們也無法實(shí)現(xiàn)異步分布式學(xué)習(xí)，并且我們需要在遷移現(xiàn)有的 tensorflow 程序上花大功夫。

TensorFlowOnSpark

我們的新框架，TensorFlowOnSpark(TFoS)，支持 TensorFlow 在 Spark 和 Hadoop 上的分布式運(yùn)行。如上圖(圖 2)所示，TFoS 與 SparkSQL、MLlib 以及其他的 Spark 庫(kù)一起在一個(gè)項(xiàng)目或線程(pipeline)中運(yùn)行。

TFoS 支持所有類型的 TensorFlow 程序，能實(shí)現(xiàn)同步和異步的訓(xùn)練與推理。并且支持模型和數(shù)據(jù)的平行處理，以及 TensorFlow 工具(如 TensorBoard)在 Spark 群集上使用。

任何 TensorFlow 程序都能夠很容易通過修改實(shí)現(xiàn)在 TFoS 上運(yùn)行的。通常情況下，只需要修改少于 10 行的 Python 代碼。很多在雅虎平臺(tái)上使用 TensorFlow 的開發(fā)者，已經(jīng)輕松將 TensorFlow 項(xiàng)目轉(zhuǎn)移到 TFoS 上執(zhí)行了。

TFoS 支持張量(tensor)在 TensorFlow 處理過程中(計(jì)算節(jié)點(diǎn)和參數(shù)服務(wù)節(jié)點(diǎn))信息的直接溝通。過程到過程(Process-to-process)的直接溝通機(jī)制使 TFoS 項(xiàng)目很容易在增加的機(jī)器上進(jìn)行擴(kuò)展。如圖 3 所示，TFoS 不需要 Spark 驅(qū)動(dòng)器(driver)參與到張量溝通中來，因此也就與具備類似于獨(dú)立 TensorFlow 群集的擴(kuò)展能力。

TFoS 提供兩種不同模式來「吞入」用于訓(xùn)練和推理的數(shù)據(jù) ：

1. TensorFlow QueueRunners：TFoS 利用 TensorFlow 的文件讀取(file readers)和 QueueRunners 來直接從 HDFS 文件中讀入數(shù)據(jù)。在數(shù)據(jù)獲取過程中不需要 Spark 參與。

2. Spark 供給：Spark RDD 數(shù)據(jù)將會(huì)被傳輸至每一個(gè) Spark 執(zhí)行器里，Spark 執(zhí)行器會(huì)進(jìn)一步將數(shù)據(jù)傳入 TensorFlow 圖(通過 feed_dict 參數(shù))。

圖 4 展示了 Inception 圖像分類網(wǎng)絡(luò)中同時(shí)進(jìn)行的分布式訓(xùn)練如何在 TFoS 中通過 QueueRunners 的一個(gè)簡(jiǎn)單設(shè)置進(jìn)行擴(kuò)展：將每個(gè)計(jì)算節(jié)點(diǎn)設(shè)置為 1 個(gè) GPU，一個(gè)讀入(reader)以及批處理數(shù)為 32。四個(gè) TFoS 的任務(wù)同時(shí)進(jìn)行以用于訓(xùn)練 10 萬步。兩天多后，當(dāng)這些任務(wù)完成時(shí)，top-5 精確度(accuracy)分別為 0.730, 0.814, 0.854，0.879。0.730 的精確度需要單計(jì)算節(jié)點(diǎn)運(yùn)行 46 小時(shí)得到，雙計(jì)算節(jié)點(diǎn)需要 22.5 個(gè)小時(shí)，4 計(jì)算機(jī)點(diǎn)需要 13 小時(shí)，8 計(jì)算節(jié)點(diǎn)需要 7.5 個(gè)小時(shí)。在 Inception 模型訓(xùn)練上，TFoS 幾乎能達(dá)到線性擴(kuò)展。這是很鼓舞人心的，雖然 TFoS 在不同模型和超參數(shù)上的擴(kuò)展能力不同。

用于分布式 TensorFlow 的 RDMA

在雅虎的 Hadoop 集群上，GPU 節(jié)點(diǎn)通過以太網(wǎng)和無線寬帶相互連接。無線寬帶提供了高速的連接，并支持在 RDMA 中直接訪問其他服務(wù)器的存儲(chǔ)。然而目前 TensorFlow 僅支持在以太網(wǎng)上使用 「gRPC」 的分布式學(xué)習(xí)。為了加速分布式學(xué)習(xí)，我們?cè)鰪?qiáng)了 TensorFlowC++層，實(shí)現(xiàn)了無線寬帶上的 RDMA。

為了結(jié)合我們發(fā)布的 TFoS，我們?cè)?default「gRPC」協(xié)議之外，引進(jìn)了一個(gè)新的 TensorFlow 服務(wù)器協(xié)議。任何分布式 tensorflow 程序都能通過指定 protocol=「grpc_rdma」in tf.train.ServerDef()or tf.train.Server() 來使用我們的增強(qiáng)版的 TensorFlow。

有了這個(gè)新協(xié)議后，就需要一個(gè) RDMA 匯集管理器(rendezvous manager)來確保張量直接寫入遠(yuǎn)程服務(wù)器的內(nèi)存。我們最大限度地減少?gòu)埩烤彌_的創(chuàng)建：張量緩沖器在開始時(shí)分配一次，然后在一個(gè) TensorFlow 工作任務(wù)的所用訓(xùn)練步驟中重復(fù)使用。從我們?cè)缙诘拇笮湍Ｐ蛯?shí)驗(yàn)，比如 VGG-19 開始，我們的就已經(jīng)證明了，與現(xiàn)有的 gRPC 相比，我們的 RDMA 實(shí)現(xiàn)在訓(xùn)練時(shí)間上帶來了顯著的提速。

由于 RDMA 支持對(duì)性能要求很高(見 TensorFlow issue#2916)，我們決定讓我們現(xiàn)有的實(shí)現(xiàn)版本作為一個(gè)預(yù)覽版向 TensorFlow 社區(qū)開放。在未來的幾周內(nèi)，我們將會(huì)進(jìn)一步優(yōu)化我們的 RDMA 實(shí)現(xiàn)，并分享一些基準(zhǔn)結(jié)果細(xì)節(jié)。

簡(jiǎn)單的 CLI 和 API

TFoS 程序是通過標(biāo)準(zhǔn)的 ApacheSpark 命令 spark-submit 運(yùn)行的。如下所示，用戶可以在 CLI 中指定 Spark 執(zhí)行器的數(shù)目，每個(gè)執(zhí)行器所用的 GPU 數(shù)目以及參數(shù)服務(wù)節(jié)點(diǎn)數(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í)行器構(gòu)建一個(gè) TensorFlow 群集
• TFCluster.start()... 在執(zhí)行器上加載 TensorFlow 程序
• TFCluster.train() or TFCluster.inference() …將 RDD 數(shù)據(jù)傳入 TensorFlow 處理
• TFCluster.shutdown() …在執(zhí)行器中結(jié)束 TensorFlow 的運(yùn)行

開源

• TensorFlowOnSpark 開源地址： github.com/yahoo/TensorFlowOnSpark
• RDMA 增強(qiáng)版開源地址： github.com/yahoo/tensorflow/tree/yahoo
• 提供多示例程序(包括 MNIST，Cifar10，Inception，and VGG)以說明 TensorFlow 程序到TensorFlowOnspar轉(zhuǎn)換過程，并且利用 RDMA。地址：https://github.com/yahoo/TensorFlowOnSpark/tree/master/examples
• 提供一張亞馬遜機(jī)器圖像用于在 AWS EC2 上應(yīng)用 TensorFlowOnSpark。接著，與 CaffeOnSpark 一樣，我們會(huì)推進(jìn) TensorFlowOnSpark。地址：https://github.com/yahoo/TensorFlowOnSpark/wiki/GetStarted_EC2

【本文是51CTO專欄機(jī)構(gòu)機(jī)器之心的原創(chuàng)文章，微信公眾號(hào)“機(jī)器之心( id: almosthuman2014)”】

戳這里，看該作者更多好文