TensorFlow On Spark 开源项目分析

原文：http://sparkdata.org/?p=423&utm_source=tuicool&utm_medium=referral 作者：京东大数据技术保障团队

概述

自Google发布TensorFlow并宣布开源，促使更多的开发者与学术研究者开始关注深度学习。各大公司也加入到深度学习的应用和开发中。TensorFlow 是谷歌基于DistBelief进行研发的第二代人工学习系统，最初由Google大脑小组开发出来，用于机器学习和深度神经网络方面的研究。这个系统的通用性使它也可以用于其他计算领域。目前大多应用于语音识别或图像识别等领域。

TensorFlow是一个采用数据流图（data flow graphs），用于数值计算的开源软件库。图中的节点（Nodes）表示数学操作，图中的线则表示节点间的相互联系的多维数据组，即张量(Tensor)。利用这个库我们可以在多种平台上展开数据分析与计算，如CPU(或GPU)， 台式机，服务器，甚至移动设备等等。

尽管TensorFlow也开放了自己的分布式运行框架，但在目前公司的技术架构和使用环境上不是那么的友好，如何将TensorFlow 加入到现有的环境中（Spark /YARN），并为用户提供更加方便易用的环境成为了目前所要解决的问题。基于以上的考虑，我们分析了目前行业内的几种TensorFlow与现有大数据计算环境结合的解决方案。目前比较流行的是TensorFlow On Spark 的解决方案，利用Spark本身的技术特性与分布式的优势使TensorFlow 并行起来。还有Google基于Docker的kerbernets方案，使TensorFlow可以在容器上轻松部署与运行。

分析内容

本次主要对TensorSpark项目，databricks 公司的TensorFrames项目，还有Yahoo公司开源的TensorFlowOnSpark项目进行分析和对比。

• TensorFlowOnSpark分析

TensorFlowOnSpark 项目是由Yahoo开源的一个软件包，能将TensorFlow与Spark结合在一起使用，为Apache Hadoop和Apache Spark集群带来可扩展的深度学习功能。使Spark能够利用TensorFlow拥有深度学习和GPU加速计算的能力。传统情况下处理数据需要跨集群（深度学习集群和Hadoop/Spark集群）,Yahoo为了解决跨集群传递数据的问题开发了TensorFlowOnSpark项目。TensorFlowOnSpark目前被用于雅虎私有云中的Hadoop集群，主要进行大规模分布式深度学习使用。

TensorFlowOnSpark在设计时充分考虑了Spark本身的特性和TensorFlow的运行机制，大大保证了两者的兼容性，使得可以通过较少的修改来运行已经存在的TensorFlow程序。在独立的TFOnSpark程序中能够与 SparkSQL、MLlib和其他 Spark 库一起工作处理数据。

图1：TensorFlowOnSpark 与现有库关系

TensorFlowOnSpark在内部实现了与Tensorflow集群类似的可扩展性。在下图中可以看出，Spark Driver端程序并不会参与TensorFlow内部相关的计算和处理，其设计思路像是一个TensorFlow集群运行在了Spark上，在每个Spark executor中启动TensorFlow组件，然后通过gRPC或 RDMA 方式进行数据传递与交互。

图2：TensorFlowOnSpark 架构

TensorFlowOnSpark程序运行时会通过如下步骤创建并管理TensorFlow集群（Spark程序会创建一个TensorFlow集群，TensorFlow的相关组件运行在Spark executor内）：

• Reservation – 为每个TensorFlow executor预留端口，并启动“数据/控制消息” 的监听程序。
• Startup – 启动TensorFlow的主程序
• Data ingestion
• Readers & QueueRunners – 利用TensorFlow的Reader机制直接从HDFS读取数据文件。
• Feeding -使用Hadoop InputFormat/OutputFormat 访问HDFS上的TFRecords数据，然后利用feed_dict机制将Spark RDD数据发送到TensorFlow节点。
• Shutdown -关闭executor上的Tensorflow worker上的线程和Parameter Server节点。

官方给出TensorFlowOnSpark目前支持的特性如下：

• 轻松迁移所有现有的TensorFlow程序，修改代码小于10行;
• 支持所有TensorFlow功能：同步/异步训练，模型/数据并行，inferencing和TensorBoard;
• Server端到Server端直接通信提升处理速度（RDMA启用时）;
• 允许由Spark分发或由TensorFlow拉取HDFS和其他资源上的数据;
• 轻松整合现有的数据处理流程和机器学习算法（如MLlib，CaffeOnSpark）;
• 轻松在云或内部部署：CPU & GPU, Ethernet and Infiniband.

我们可以使用两种方法来提取训练数据和预测：

• TensorFlow QueueRunners：TensorFlowOnSpark利用TensorFlow的文件读取器和QueueRunners直接从HDFS文件读取数据。 Spark不涉及访问数据。
• Spark Feeding：通过Spark RDD将数据传输到每个Spark executor，然后通过feed_dict将数据传输到TensorFlow graph中。

开发的TFoS （TensorFlowOnSpark）程序可以直接使用Spark的Spark-submit命令提交到集群上，在提交时程序时，用户可以指定Spark executor的个数，每个executor上有几个GPU，“参数服务器（Parameter Server）”的个数。另外用户还可以指定TensorBoard和RDMA是否使用, 命令如下：

sparksubmit –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：

• reserve() … construct a TensorFlow cluster from Sparkexecutors
• start() … launch Tensorflow program on the executors
• train() or TFCluster.inference() … feed RDD data to TensorFlowprocesses
• TFCluster.shutdown() … shutdown Tensorflow execution on executors

优点：

• Yahoo已经在自己的云上实践并使用，可以用于生产
• 轻松整合现有的数据处理流程和机器学习算法
• 支持所有TensorFlow功能
• 轻松移植现有TensorFlow程序到Spark集群上
• 学习成本较低

缺点：

• 缺少相关文档与使用教程
• 代码刚开源，社区反馈与支持力度未知
• TensorFrames 分析

TensorFrames 是databricks公司开发的一个软件包，利用它可以使用类似Spark DataFrames 的方式操作TensorFlow程序。官方描述目前这个软件包是实验性的，仅作为技术预览提供。部分功能还存在性能问题。有了它我们可以方便的使用我们熟悉的Spark 开发环境进行Tensorflow 深度学习应用开发，大大降低了学习成本。官方给出了一个简单的使用案例，从这个案例中可以看出，其使用方法和传统的DataFrames方式的编程非常相似。下面的这个案例中实现了一个简单的分布式打印”Hello TensorFlow”的应用。

图3： Hello word案例代码

官方文章从两个角度介绍了使用这个软件包能够解决的问题：

1.超参数调优：使用Spark找到神经网络训练的最佳超参数集，可以将训练时间减少10倍，误差率降低34％。

2.大规模部署模型：使用Spark对大量数据应用经过训练的神经网络模型。

从原理上TensorFlow库会自动创建各种形状和大小的神经网络的训练算法。 然而，构建神经网络的实际过程比在数据集上运行某些函数更复杂。通常需要设置许多非常重要的超参数，这些参数会影响如何训练模型。在实际操作中，机器学习开发人员会使用不同的超参数重复运行相同的模型多次，以便找到最佳值。从这个角度考虑如果能让Spark利用多台机器同时运行多组参数计算模型，那么将会对性能有很大的提升。

作者利用这个特性并行计算13个模型，相当于在一台机器上一次训练一个模型的7倍速度。

优点：

• 大公司出品，后期可能会加速推广与应用
• 较高层次封装，可以使用Spark DataFrame相关特性

缺点：

• 相对于TensorFlow集群并行的灵活性而言，这个框架提供的API较少
• 缺少相关使用文档
• 性能和功能缺失，还不能用于生产
• TensorSpark 分析

TensorSpark 是本次分析中比较简单的一个项目，这个项目在GitHub上开源。作者在其博客中写到在Spark上运行TensorFlow的基本问题是如何在Spark上分配神经网络的训练。 Spark对于迭代map-reduce问题非常有用，但是训练神经网络不是一个map-reduce问题。 作者受到谷歌的SGD架构启发。如果模型能够在单个节点装下，那么采用数据并行是最适合的。利用SGD（既DistBelief）结构的参数服务器可以很容易实现分布式深度学习，这种方式很适合于大量数据和较小模型的训练方式。

上图左侧是论文中的SGD架构处理流程，右侧为用Python实现了简单的参数服务器结构的Spark程序。

工作流程是Spark Worker启动单机版的Tensorflow异步计算梯度，周期性把梯度发给Driver，就是参数服务器，在汇总多个Worker的梯度后，参数服务器把梯度数据返回给Worker。这种方式的并行实现简单，但Driver（参数服务器）很容易成为瓶颈影响扩展。

使用方法为：

• 将程序代码打包成zip包
• 使用spark-submit 命令将zip包，提交到集群执行

官方例子如下：

zip pyfiles.zip ./parameterwebsocketclient.py ./parameterservermodel.py ./mnistcnn.py ./mnistdnn.py ./moleculardnn.py ./higgsdnn.py

spark-submit \ –master yarn \ –deploy-mode cluster \ –queue default \ –num-executors 3 \ –driver-memory 20g \ –executor-memory 60g \ –executor-cores 8 \ –py-files ./pyfiles.zip \ ./tensorspark.py

脚本功能定义：

tensorspark/gpu_install.sh – script to build tf from source with gpu support for aws tensorspark/simple_websocket_.py – simple tornado websocket example tensorspark/parameterservermodel.py – “abstract” model class that has all tensorspark required methods implemented tensorspark/dnn.py – specific fully connected models for specific datasets tensorspark/mnistcnn.py – convolutional model for mnist tensorspark/parameterwebsocketclient.py – spark worker code tensorspark/tensorspark.py – entry point and spark driver code

优点：

• 框架代码量少，学习成本较低
• 安装方便

缺点：

• 需要自己基于Spark框架实现参数服务器和通讯逻辑
• 可扩展性较低，目前框架不能充分利用tensorflow 本身的其他特性
• 现在处于科研性质，不太适用于生产推广

结论

通过对这三个项目的分析，可以看出开源社区有很多家公司都在积极的推动TensorFlow 在Spark上运行实验和开发。每个项目都有自己的特点和特定使用方式。

TensorSpark项目的确是简陋了一点，而且貌似不再更新了，但这个项目启发了我们如何设计和开发TensorFlow on YARN 模型的一种思路。

TensorFrames这个项目对TensorFlow做了高度抽象的封装，将TensorFlow做成了Spark的一个底层，如果想要将已有的TensorFlow程序运行在这个框架上，那么需要修改代码（修改量视复杂程序算）。

对于Yahoo的TensorFlowOnSpark项目，从官方介绍上看是已经投入生产并与现有TensorFlow程序能够最大兼容的一个框架。现有的TensorFlow程序可以做较少的修改运行在这个框架上，而且支持所有TensorFlow的特性和功能。

本文是基于官方的介绍整理归纳而成，由于水平有限不免有错误之处，欢迎指正。