maven依赖 spark sql_社区分享 | Spark 玩转 TensorFlow 2.0

本文来自社区投稿与征集，作者梁云，转自：https://github.com/lyhue1991/eat_tensorflow2_in_30_days

本篇文章介绍在 Spark 中调用训练好的 TensorFlow 模型进行预测的方法。本文内容的学习需要一定的 Spark 和 Scala 基础。如果使用 Pyspark 的话会比较简单，只需要在每个 executor 上用 Python 加载模型分别预测就可以了。但工程上为了性能考虑，通常使用的是 Scala 版本的 Spark。本篇文章我们通过 TensorFlow for Java 在 Spark 中调用训练好的 TensorFlow 模型。利用 Spark 的分布式计算能力，从而可以让训练好的 TensorFlow 模型在成百上千的机器上分布式并行执行模型推断。本案例以 TensorFlow 2.0 的  tf.keras 接口训练的线性模型为例进行演示。在本例基础上稍作修改则可以用 Spark 调用训练好的各种复杂的神经网络模型进行分布式模型推断。但实际上 TensorFlow 并不仅仅适合实现神经网络，其底层的计算图语言可以表达各种数值计算过程。利用其丰富的低阶 API，我们可以在 TensorFlow 2.0 上实现任意机器学习模型。结合  tf.Module 提供的便捷的封装功能，我们可以将训练好的模型导出成模型文件并在 Spark 上分布式调用执行。

这无疑为我们的工程应用提供了巨大的想象空间。

Spark-Scala 调用 TensorFlow 模型概述

在 Spark(Scala) 中调用 TensorFlow 模型进行预测需要完成以下几个步骤：

1. 准备 protobuf 模型文件

2. 创建 Spark-Scala 项目，在项目中添加 Java 版本的 TensorFlow 对应的 jar 包依赖

3. 在 Spark-Scala 项目中 driver 端加载 TensorFlow 模型调试成功

4. 在 Spark-Scala) 项目中通过 RDD 在 executor 上加载 TensorFlow 模型调试成功

5. 在 Spark-Scala 项目中通过 DataFrame 在 executor 上加载 TensorFlow 模型调试成功

** 建议参考本项目中的《Appendix I: Load TensorFlow model in Spark.ipynb》文档。

一 准备 protobuf 模型文件

我们使用 tf.keras 训练一个简单的线性回归模型，并保存成 protobuf 文件。

import tensorflow as tf
from tensorflow.keras import models,layers,optimizers

## 样本数量
n = 800

## 生成测试用数据集
X = tf.random.uniform([n,2],minval=-10,maxval=10)
w0 = tf.constant([[2.0],[-1.0]])
b0 = tf.constant(3.0)

Y = X@w0 + b0 + tf.random.normal([n,1],mean = 0.0,stddev= 2.0)  # @表示矩阵乘法,增加正态扰动

## 建立模型
tf.keras.backend.clear_session()
inpu