雅虎开源了 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
感谢冬雨对本文的审校。
给InfoQ 中文站投稿或者参与内容翻译工作,请邮件至 editors@cn.infoq.com 。也欢迎大家通过新浪微博( @InfoQ , @丁晓昀),微信(微信号: InfoQChina )关注我们。
评论