Tensorflow for Java + Spark-Scala 分布式机器学习计算框架的应用实践

王辉，2017 年加入去哪儿网。

目前负责反爬虫相关风控业务，技术领域涉猎广泛，在风控智能化实践方向的道路上持续探索中。

一. 前言

Qunar 智能风控场景中，风控研发团队经常会应用一些算法模型，来解决复杂场景问题。典型的如神经网络模型，决策树模型等等。而要完成模型从训练到部署预测的全过程，除了模型算法之外，离不开技术框架的支撑。本篇文章将和大家分享一下，在预测服务部署阶段，基于 Tensorflow for Java 和 Spark-Scala 构建分布式机器学习计算框架的实践经验。主要围绕以下几点展开：

• Tensorflow for Java  &  Spark-Scala 是什么？

• 框架选型和适用场景思考

• 如何使用 Spark-Scala 集成 Tensorflow for Java 构建预测服务?

• 框架实践过程中的优化以及踩坑经验

二. 框架选型

2.1 项目场景

项目的背景，是对手机客户端收集的用户风控数据采集后，构建神经网络模型，离线分析预测用户风险。

客户端收集的数据量很大，目前阶段小时特征数据量约 300w ，后续还会扩大到千万级别。同时，我们希望每小时执行一次预测任务，那么需要保证一次计算能在小时内完成，否则会出现任务积压。所以即使是离线计算这种实时性要求不高的场景，也依然有着高性能处理的要求。

在这个背景下，我们期望基于大数据分布式机器学习计算框架来提升模型预测效率。

2.2 框架选型

在分布式机器学习计算框架选型上，一类是机器学习框架自身支持的分布式能力，比如 Tensorflow、Pytorch 两大主流框架都已支持分布式计算；另一类是大数据分布式框架和机器学习框架结合，通过 Spark、Flink 等大数据框架集成机器学习框架 API。

目前虽然机器学习框架逐渐具备分布式能力，但主要目的是解决数据量大、模型参数多场景下的模型训练性能问题。相对于传统大数据分布式框架，优点是对于分布式训练的支持和针对机器学习场景的分布式方案。缺点是在数据归并和多级分层上没有提供很好的解决方案，其次，需要单独搭建和维护集群，增加了运维成本。综合考虑，在分布式预测场景下，采用了大数据框架和机器学习框架结合的方案。

大数据框架通常分为流处理和批处理两类。流处理适用于实时和准实时计算， Flink 、 Storm 、 Spark-Streaming 等都属于流处理框架；批处理适用于离线计算，我们的场景是典型的离线批处理场景，主流的批处理框架目前有 Spark 和 Hadoop 的 MapReduce 。通过框架实现上的对比，可以发现 Spark 基于内存运算虽然消耗更多资源，但在性能上能带来很大的提升。

而在机器学习框架领域， TensorFlow 、 PyTorch 目前分别成为了工业界和学术界使用最广泛的两大框架。TensorFlow 是谷歌的开发者创造的一款开源的深度学习框架，于 2015 年发布。PyTorch 是最新的深度学习框架之一，由 Facebook 的团队开发，并于 2017 年在 GitHub 上开源。

我们从多个维度对比下 Tensorflow 和 Pytorch 两种框架，可以发现：Pytorch 具备更优秀的 API ，上手快，动态图，易调试等优点，适合研究者快速构建模型，迭代速度快。Tensorflow 在多语言支持、跨平台能力、性能、生产部署等方面具备优势，适合于生产部署。这也是为什么 Pytorch 更受学术界欢迎，而 Tensorflow 保持着工业界的主导地位。Tensorflow 和 Pytorch 目前都在试图向对方的优势靠拢，不过短期内这样的现状应该不会改变。

通过上面的分析，结合实际应用场景，风控模型的算法结构并不是很复杂，在 Tensorflow 和 Pytorch 下都能快速构建，因此考虑到模型生产环境部署难度，以及跨平台部署的需求，我们选择 Tensorflow 来构建模型。

2.3 为什么选择 Tensorflow for Java & Spark - Scale

在 Spark 和 Tensorflow 集成的方案中，一种是在 Python 环境中使用 PySpark+Tensorflow for Python 来实现，这也是目前多数的实践方案。另一种就是 Java/Scala 环境中 Spark-Scala+Tensorflow for Java 的方案。我们选择第二种方案，主要是出于性能考虑，下文将对比两种框架底层的性能差异：

2.3.1 Tensorflow for Java & Spark - Scale 介绍

TensorFlow 已支持多种客户端语言下的安装和运行，不过 Python 仍是目前唯一受到良好支持的语言。Tensorflow for Java 是 TensorFlow 提供的用于 Java 程序的 API ，这些 API 适合于加载在 Python 中创建的模型并在 Java 应用程序中执行。

Spark-Scala 也就是 Scala 语言环境下的 Spark 任务开发工具。因为 Scala 语言也是运行在 JVM 上，且兼容 Java 语言，所以可以集成 Java 应用。Spark 框架的开发语言也是 Scala ，这是因为 Scala 采用函数式编程在并行和并发计算编程方面具备优势。

2.3.2 性能对比

上面我们说了选择 Tensorflow for Java & Spark-Scala 的目的是满足高性能的场景需求。下面分析下两种方案性能上的差异。

Spark 的运行时架构如下。用户的 Spark 应用程序运行在 Driver 上，经过 Spark 调度封装成一个个 Task ，再将这些 Task 信息发给 Executor 执行。Scala 版本的 Spark 运行在这种原生架构下。

PySpark 的运行时结构如下。为了不破坏 Spark 已有的运行时架构，Spark 在外围包装一层 Python API ，借助 Py4j 实现 Python 和 Java 的交互，进而实现通过 Python 编写 Spark 应用程序。

由此可以得知， PySpark 在性能上弱于 Spark-Scala ，主要原因有两点：

• 第一点， Driver 端多了一层 Python 到 Java 的转换；

• 第二点， Executor 端为了运行用户定义的 Python 函数或 Lambda 表达式，为每个 Task 单独启一个 Python 进程，通过 socket 通信方式将 Python 函数或 Lambda 表达式发给 Python 进程执行。

再对比下 Tensorflow for Python 和 Tensorflow for Java 的差异。两者底层都是 Tensorflow 的 C++ 函数库，性能上差异不大。上层来说 Java 语言会比 Python 更快些，对于调用 Tensorflow API 之前需要做许多预处理操作的场景会体现出差异。

综上得出结论，方案二在性能上优于方案一。

当然，方案一在其他方面也有着其优势，比如开发效率高、集成难度低(无需跨平台)、 API 支持度高等等。

2.4 小节

本节内容总结：

• 基于 Spark 大数据框架和 Tensorflow 机器学习框架结合，实现分布式机器学习预测，是一种相对可行、有效的方案。

• 基于 Tensorflow for Java 和 Spark-Scala 实现 Spark 和 Tensorflow 集成，能带来更高的性能。

• 当前方案适用于大数据下高性能分布式机器学习模型离线预测场景。

三. 应用实践

实践中项目应用流程：首先基于样本数据训练好模型，生成模型文件。之后在 Spark 中读取特征数据，调用 Tensorflow Java API 加载模型，进行预测，得到结果集。

3.1 模型训练

3.1.1 训练模型文件

在我们的项目中模型训练基于 Python+Tensorflow+Keras 实现。这里以 MNIST 数据集 CNN 分类为例演示模型训练代码。经过训练后得到了保存为 protobuf 格式的模型文件， protobuf 格式文件可以跨平台加载出模型。

import tensorflow as tffrom tensorflow.keras.models import Sequentialfrom tensorflow.keras.layers import Dense,Dropout,Convolution2D,MaxPooling2D,Flattenfrom tensorflow.keras.optimizers import Adam
def train_model():    # 载入训练集和测试集数据，进行独热编码    mnist = tf.keras.datasets.mnist    (x_train, y_train), (x_test, y_test) = mnist.load_data()      y_train = tf.keras.utils.to_categorical(y_train,num_classes=10)      y_test = tf.keras.utils.to_categorical(y_test,num_classes=10)      # 定义顺序模型      model = Sequential()      # 卷积层、池化层、扁平化、全连接      model.add(Convolution2D(input_shape=(28, 28, 1), filters=32, kernel_size=5, strides=1, padding='same', activation='relu'))      model.add(MaxPooling2D(pool_size=2, strides=2, padding = 'same'))      model.add(Convolution2D(64, 5, strides=1, padding='same', activation='relu'))      model.add(MaxPooling2D(2,2,'same'))      model.add(Flatten())      model.add(Dense(1024,activation = 'relu'))      model.add(Dropout(0.5))      model.add(Dense(10,activation='softmax'))      # 定义优化器，loss function，训练过程中计算准确率      adam = Adam(lr=1e-4)      model.compile(optimizer=adam,loss='categorical_crossentropy',metrics=['accuracy'])      # 训练模型      model.fit(x_train,y_train,batch_size=64,epochs=10,validation_data=(x_test, y_test))      # 保存模型      model.save('./model/model_v1', save_format="tf")

3.1.2 查看模型文件

进入模型文件目录，执行以下命令，可以展示模型文件信息。圈红的信息由上到下依次为模型的标签，签名，输入张量，输出张量，预测方法名。在之后加载模型预测时会用到这些信息。

saved_model_cli show --dir ./model_v1/ --all

3.2 模型预测

3.2.1 工程搭建 &框架引入

新建 Scala 工程，引入 Spark 和 Tensorflow 依赖

<!-- scala --><dependency>    <groupId>org.scala-lang</groupId>   		<artifactId>scala-library</artifactId>   		<version>${scala.version}</version></dependency><!-- spark hadoop --><dependency>   		<groupId>org.apache.spark</groupId>   		<artifactId>spark-core_${spark.scala.version}</artifactId>   		<version>${spark.version}</version></dependency><dependency>    		<groupId>org.apache.spark</groupId>    		<artifactId>spark-hive_${spark.scala.version}</artifactId>   		<version>${spark.version}</version></dependency><dependency>    		<groupId>org.apache.spark</groupId>   		<artifactId>spark-sql_${spark.scala.version}</artifactId>    		<version>${spark.version}</version></dependency><!-- tensorflow --><dependency>    		<groupId>org.tensorflow</groupId>    		<artifactId>tensorflow</artifactId>        <version>1.15.0</version></dependency>

3.2.2 模型文件加载

调用 Tensorflow API 加载预训练好的 protobuff 格式模型文件，得到 SavedModelBundle 类型模型对象。模型文件我们可以保存在工程 resource 目录下，再从 resource 目录加载（ Tensorflow 不支持直接从 HDFS 记载模型，后文会介绍如何实现）。

package com.tfsparkimport org.apache.spark.sql.SparkSessionimport org.tensorflow.SavedModelBundleimport org.{tensorflow => tf}
object ModelLoader {  	//modelPath是模型在resource下路径，modelTag从模型文件信息中获取  	def loadModelFromLocal(spark: SparkSession, modelPath: String, modelTag: String): SavedModelBundle = {    		val bundle = tf.SavedModelBundle.load(modelPath, modelTag)    }  }

3.2.3 调用 Tensorflow API 预测

在 Java 版本的 Tensorflow 中还是类似 Tensorflow1.0 中静态计算图的模式，需要建立 session ，指定 feed 的特征数据和 fetch 的预测结果，然后执行 run 方法。

查看模型文件获取的信息将在这里作为参数传入。

package com.tfspark.tensorflowimport com.qunar.rdc.util.TfUtilimport org.tensorflow.SavedModelBundleimport scala.collection.mutable.WrappedArrayimport org.{tensorflow => tf}
object TensorFlowCnnProcessor {
	def predict(broads: SavedModelBundle, features: WrappedArray[WrappedArray[WrappedArray[Float]]]): Int = {		val sess = bundle.session()		// 特征数据格式化		val x = tf.Tensor.create(Array(features.map(a => a.map(b => b.toArray).toArray).toArray))		// 执行预测 需要传入模型信息里的输入张量名和输出张量名，以及格式化后的特征数据		val y =  sess.runner().feed("serving_default_hmc_input:0", x).fetch("StatefulPartitionedCall:0").run().get(0)		// 结果是1x2的二维数组		val result = Array.ofDimFloaty.copyTo(result)    // 返回最大值坐标，即为分类结果，对应的是one-hot编码		TfUtil.argMaxOneDim(result(0))	}
}

3.2.4 Spark 结合 Tensorflow 预测

Spark 从 Hive 读取预测数据，经过预处理转换成特征数据，调用 Tensorflow API 预测。通过 Tensorflow API 与 Spark 分布式数据集结合使用，实现分布式批处理框架和机器学习的集成。

// 将封装Tensorflow API的预测方法注册为udf函数val sensorPredict = udf((features: WrappedArray[WrappedArray[WrappedArray[Float]]]) => {predict(bundle, features)})// Dataframe调udf函数val resultDf = featureDf.withColumn("predict_result", sensorPredict(col("feature"))

3.3 服务部署

3.3.2 环境依赖

将 Spark-Scala 和 Tensorflow for Java 集成后的工程，通过 maven 打出依赖包：tfspark-1.0.0-jar-with-dependencies.jar 。

在部署了 spark 运行环境的 hadoop 集群上运行 jar 包。依赖的集群环境需提前安装 spark、hadoop、hive 等大数据组件。

3.3.2 执行脚本

spark-submit 执行 jar 包，指定执行的 main 函数类 com.tfspark.PredictMain ，指定 jar 包路径，设置执行任务的 executor 数和核心数以及内存参数，传入模型文件版本参数。

sudo -u root /usr/local/Cellar/apache-spark/2.4.3/bin/spark-submit --class com.tfspark.PredictMain --master yarn --deploy-mode client --driver-memory 6g --executor-memory 6g --num-executors 5 --executor-cores 4  /tmp/tfspark-1.0.0-jar-with-dependencies.jar model_v1

3.4 实践成果

完成 Tensorflow for Java 和 Spark-Scala 的集成，实现大数据分布式批处理框架和机器学习的结合。将 Python 环境下生成的模型文件，加载应用于 Java 平台， 达到机器学习模型跨平台应用的效果。

顺利应用于线上项目，每小时完成 300w 数据模型预测，任务耗时 9m ，吞吐量达到 5500+/秒。实现大数据场景下高性能的离线模型预测，打通了整套应用流程。

四. 优化 &踩坑经验

市场性能优化

在 3.2.4 节示范了 Spark 在 DataFrame 中调用 Tensorflow API 的常规操作流程。我们的项目按以上实现方式上线之初， 300w 数据执行耗时在 20m 左右。分析之后认为性能上有优化的空间。

• 问题点：每一条数据都会调用一次模型预测方法，会导致一些可复用的对象被多次创建，相同的方法流程也被多次调用。

• 优化思路：数据批量调用预测方法。减少重复的对象创建和方法流程执行。

• 解决方案：使用 RDD 模式下 mapPartition 算子替代 map 算子，获取特征数组，批量调用。

对比下 mapPartition 算子和 map 算子的实现：

两者都是操作 partition 的迭代器， map 算子通过迭代器获取每个元素，调用操作函数，函数入参是元素类型。mapPartition 直接将迭代器传给操作函数，函数入参是元素集合的迭代器类型。所以使用区别在于， mapPartition 在一个方法中，操作所有 partition 元素，调用一次操作函数；map 一次只能操作一个元素，调用多次操作函数。

因此 mapPartition 对比 map ，更适用于存在重复对象创建或流程调用的场景，可以提升性能效率；mapPartition 存在的突出缺点是可能导致 OOM ，因为一次加载多个元素，相对于 map 一次加载一个元素，占用内存更多，不能及时垃圾回收。

Tensorflow API 支持传入数组批量调用，通过 mapPartition 将迭代器转换成数组，就可以批量预测，提升了效率。

float[][] matrix = new float[m][n];Tensor<Float> ft = Tensor.create(matrix, Float.class);val y =  sess.runner().feed("serving_default_hmc_input:0", ft).fetch("StatefulPartitionedCall:0").run().get(0)

结果：经过采用 RDD 模式下 mapPartition 算子实现批量预测后，任务时长显著下降，由 20m 降至 9m。

模型文件热更新

• 问题点：上面我们提到模型文件存放在工程 resource 目录，在模型结构不变的情况下，这种方式不便于对模型文件进行更新，需要重新部署服务。

• 解决方案：对存储方式进行改进，将模型文件存放在 HDFS 中。每次从 HDFS 获取模型数据。Tensorflow 本身没有提供直接从 HDFS 加载模型的 API，但可以通过 Spark 先从 HDFS 读到本地，再从本地加载来实现。这样在模型结构不变的情况下，每次只需要上传新的模型文件，对 HDFS 中原文件进行覆盖或升级版本号，就可以热更新。

//modelPath是模型在HDFS下路径，modelTag从模型文件信息中获取spark.sparkContext.addFile(modelPath, true)val localPath = SparkFiles.get(modelPath)tf.SavedModelBundle.load(localPath, modelTag)

Tensorflow 库文件找不到

报错：在 hadoop 集群上执行预测任务时，报错 libtensorflow_jni.so 文件找不到。

Exception in thread "main" java.lang.UnsatisfiedLinkError: /tmp/tensorflow_native_libraries-1613705012956-0/libtensorflow_jni.so: libtensorflow_framework.so.1: cannot open shared object file: No such file or directory    at java.lang.ClassLoader$NativeLibrary.load(Native Method)    at java.lang.ClassLoader.loadLibrary0(ClassLoader.java:1941)    at java.lang.ClassLoader.loadLibrary(ClassLoader.java:1824)    at java.lang.Runtime.load0(Runtime.java:809)    at java.lang.System.load(System.java:1086)

• 原因分析：通过查看报错日志，分析得出结论，Tensorflow 依赖 C 环境，对 C 基础库的版本有要求。集群 C 环境基础库版本存在不一致，部分安装了过低或过高版本，导致 Tensorflow 不兼容，部分机器任务报错。

• 解决方案：提供两种思路。

  第一种方案是修复集群环境，但如果是公共集群更改基础库影响较大。

  第二种方案是不用 Hadoop 集群，使用单台实体机，采用 Spark Local 模式，启动多个 Executor 执行任务，从而保证环境一致性。

实践中我们采用了第二种。

五. 总结

本文写作的目的，是希望将自己在分布式机器学习计算框架实际应用中的思考和经验分享出来，供大家参考交流。

通过不同框架的优缺点对比，以及底层实现对性能影响的剖析，阐述了选型的思考过程。明确了 Tensorflow for Java & Spark-Scala 为何适用于大数据下高性能分布式机器学习模型预测场景。结合实践经验，演示了项目中框架应用的整体流程。并总结了在性能和部署流程优化过程中的思考。

由于水平有限，文章多有纰漏不足，也恳请大家指正。