雅虎開源了TensorFlowOnSpark,數據科學家和工程師們可以直接利用運行于CPU/GPU架構上的Spark或者Hadoop做分布式模型訓練。據報道,這個庫支持把現有的TensorFlow程序切換到新的API,同時實現了模型訓練的性能提升。
在開源公告里,雅虎說明了TensorFlowOnSpark想解決的問題,比如在深度神經網絡訓練中管理Spark數據管線之外的其他集群帶來的運維負載,以網絡I/O為瓶頸的數據集在訓練集群的傳入和傳出,令人討厭的系統復雜性,以及端到端的整體學習時延。TensorFlowOnSpark的工作和雅虎之前開源的CaffeOnSpark相似。現有的對TensorFlow和Spark的集成所做的努力,有DataBricks的TensorFrame,以及Amp Lab 的SparkNet,這些對于雅虎來說都是在正確方向上的邁進,但是在允許TensorFlow進程之間直接通信方面還是有所欠缺。雅虎的目標之一,是讓TensorFlowOnSpark成為一個完全對Spark兼容的API,在一個Spark處理工作流里,其集成能力能跟SparkSQL、MLib以及其他Spark核心庫一樣好。
在架構上,它把給定TensorFlow算法和TensorFlow core放在一個Spark Executor中,并讓TensorFlow任務能夠通過TensorFlow的文件閱讀器和QueueRunners直接獲取HDFS數據,這是一種有著更少網絡I/O以及“把計算帶給數據”的方案。TensorFlowOnSpark在語義上就支持對執行器的端口預留和監聽,對數據和控制函數的消息輪詢,TensorFlow主函數的啟動,數據注入,直接從HDFS讀取數據的閱讀器和queue-runner機制,通過feed_dict向TensorFlow注入Spark RDD,以及關機。
除了TensorFlowOnSpark,雅虎還在他們自己的分支上擴展了TensorFlow核心C++引擎以在Infiniband里使用RDMA,這個需求在TensorFlow主項目里被提出過還產生了相關討論。雅虎的Andy Feng注意到,使用RDMA而不是gRPC來做進程間通信,在不同的網絡里會帶來百分之十到百分之兩百不等的訓練速度的提升。
查看英文原文:Yahoo Open Sources TensorFlowOnSpark