TensorFlow On Flink 原理解析

做者：陈戊超（仲卓），阿里巴巴技术专家

深度学习技术在当代社会发挥的做用愈来愈大。目前深度学习被普遍应用于个性化推荐、商品搜索、人脸识别、机器翻译、自动驾驶等多个领域，此外还在向社会各个领域迅速渗透。

背景

当前，深度学习的应用愈来愈多样化，随之涌现出诸多优秀的计算框架。其中 TensorFlow，PyTorch，MXNeT 做为普遍使用的框架更是备受瞩目。在将深度学习应用于实际业务的过程当中，每每须要结合数据处理相关的计算框架如：模型训练以前须要对训练数据进行加工生成训练样本，模型预测过程当中须要对处理数据的一些指标进行监控等。在这样的状况下，数据处理和模型训练分别须要使用不一样的计算引擎，增长了用户使用的难度。

本文将分享如何使用一套引擎搞定机器学习全流程的解决方案。先介绍一下典型的机器学习工做流程。如图所示，整个流程包含特征工程、模型训练、离线或者是在线预测等环节。

在此过程当中，不管是特征工程、模型训练仍是模型预测，中间都会产生日志。须要先用数据处理引擎好比 Flink 对这些日志进行分析，而后进入特征工程。再使用深度学习的计算引擎 TensorFlow 进行模型训练和模型预测。当模型训练好了之后再用 tensor serving 作在线的打分。

上述流程虽然能够跑通，但也存在必定的问题，好比：

1. 同一个机器学习项目在作特征工程、模型训练、模型预测时须要用到 Flink 和 TensorFlow 两个计算引擎，部署相对而言更复杂。
2. TensorFlow 在分布式的支持上还不够友好，运行过程当中须要指定机器的 IP 地址和端口号；而实际生产过程常常是运行在一个调度系统上好比 Yarn，须要动态分配 IP 地址和端口号。
3. TensorFlow 的分布式运行缺少自动的 failover 机制。

针对以上问题，咱们经过结合 Flink 和 TensorFlow，将 TensorFlow 的程序跑在 Flink 集群上的这种方式来解决，总体流程以下：

特征工程用 Flink 去执行，模型训练和模型的准实时预测目标使 TensorFlow 计算引擎能够跑在 Flink 集群上。这样就能够用 Flink 一套计算引擎去支持模型训练和模型的预测，部署上更简单的同时也节约了资源。

Flink 计算简介

Flink 是一款开源大数据分布式计算引擎，在 Flink 里全部的计算都抽象成 operator，如上图所示，数据读取的节点叫 source operator，输出数据的节点叫 sink operator。source 和 sink 中间有多种多样的 Flink operator 去处理，上图的计算拓扑包含了三个 source 和两个 sink。

机器学习分布式拓扑

机器学习分布式运行拓扑以下图所示：

在一个机器学习的集群当中，常常会对一组节点（node）进行分组，如上图所示，一组节点能够是 worker（运行算法），也能够是 ps（更新参数）。

如何将 Flink 的 operator 结构与 Machine Learning 的 node、Application Manager 角色结合起来？下面将详细讲解 flink-ai-extended 的抽象。

Flink-ai-extended 抽象

首先，对机器学习的 cluster 进行一层抽象，命名为 ML framework，同时机器学习也包含了 ML operator。经过这两个模块，能够把 Flink 和 Machine Learning Cluster 结合起来，而且能够支持不一样的计算引擎，包括 TensorFlow。

以下图所示：

在 Flink 运行环境上，抽象了 ML Framework 和 ML Operator 模块，负责链接 Flink 和其余计算引擎。

ML Framework

ML Framework 分为 2 个角色。

1. Application Manager（如下简称 am） 角色，负责管理全部 node 的节点的生命周期。
2. node 角色，负责执行机器学习的算法程序。

在上述过程当中，还能够对 Application Manager 和 node 进行进一步的抽象，Application Manager 里面咱们单独把 state machine 的状态机作成可扩展的，这样就能够支持不一样类型的做业。

深度学习引擎，能够本身定义其状态机。从 node 的节点抽象 runner 接口，这样用户就能够根据不一样的深度学习引擎去自定义运行算法程序。

ML Operator

ML Operator 模块提供了两个接口：

1. addAMRole，这个接口的做用是在 Flink 的做业里添加一个 Application Manager 的角色。Application Manager 角色如上图所示就是机器学习集群的管理节点。
2. addRole，增长的是机器学习的一组节点。

利用 ML Operator 提供的接口，能够实现 Flink Operator 中包含一个Application Manager 及 3 组 node 的角色，这三组 node 分别叫 role a、 role b,、role c，三个不一样角色组成机器学习的一个 cluster。如上图代码所示。Flink 的 operator 与机器学习做业的 node 一一对应。

机器学习的 node 节点运行在 Flink 的 operator 里，须要进行数据交换，原理以下图所示：

Flink operator 是 java 进程，机器学习的 node 节点通常是 python 进程，java 和 python 进程经过共享内存交换数据。
TensorFlow On Flink

TensorFlow 分布式运行

TensorFlow 分布式训练通常分为 worker 和 ps 角色。worker 负责机器学习计算，ps 负责参数更新。下面将讲解 TensorFlow 如何运行在 Flink 集群中。

TensorFlow Batch 训练运行模式

Batch 模式下，样本数据能够是放在 HDFS 上的，对于 Flink 做业而言，它会起一个source 的 operator，而后 TensorFlow 的 work 角色就会启动。如上图所示，若是 worker 的角色有三个节点，那么 source 的并行度就会设为 3。同理下面 ps 角色有 2 个，因此 ps source 节点就会设为 2。而 Application Manager 和别的角色并无数据交换，因此 Application Manager 是单独的一个节点，所以它的 source 节点并行度始终为 1。这样 Flink 做业上启动了三个 worker 和两个 ps 节点，worker 和 ps 之间的通信是经过原始的 TensorFlow 的 GRPC 通信来实现的，并非走 Flink 的通讯机制。

TensorFlow stream 训练运行模式

如上图所示，前面有两个 source operator，而后接 join operator，把两份数据合并为一份数据，再加自定义处理的节点，生成样本数据。在 stream 模式下，worker 的角色是经过 UDTF 或者 flatmap 来实现的。

同时，TensorFlow worker node 有3 个，因此 flatmap 和 UDTF 相对应的 operator 的并行度也为 3， 因为ps 角色并不去读取数据，因此是经过 flink source operator 来实现。

下面咱们再讲一下，若是已经训练好的模型，如何去支持实时的预测。

使用 Python 进行预测

使用 Python 进行预测流程如图所示，若是 TensorFlow 的模型是分布式训练出来的模型，而且这个模型很是大，好比说单机放不下的状况，通常出如今推荐和搜索的场景下。那么实时预测和实时训练原理相同，惟一不一样的地方是多了一个加载模型的过程。

在预测的状况下，经过读取模型，将全部的参数加载到 ps 里面去，而后上游的数据仍是通过和训练时候同样的处理形式，数据流入到 worker 这样一个角色中去进行处理，将预测的分数再写回到 flink operator，而且发送到下游 operator。

使用 Java 进行预测

如图所示，模型单机进行预测时就不必再去起 ps 节点，单个 worker 就能够装下整个模型进行预测，尤为是使用 TensorFlow 导出 save model。同时，由于 saved model 格式包含了整个深度学习预测的所有计算逻辑和输入输出，因此不须要运行 Python 的代码就能够进行预测。

此外，还有一种方式能够进行预测。前面 source、join、UDTF 都是对数据进行加工处理变成预测模型能够识别的数据格式，在这种状况下，能够直接在 Java 进程里面经过 TensorFlow Java API，将训练好的模型 load 到内存里，这时会发现并不须要 ps 角色， worker 角色也都是 Java 进程，并非 Python 的进程，因此咱们能够直接在 Java 进程内进行预测，而且能够将预测结果继续发给 Flink 的下游。

总结

在本文中，咱们讲解了 flink-ai-extended 原理，以及Flink 结合 TensorFlow 如何进行模型训练和预测。但愿经过本文大分享，你们可以使用 flink-ai-extended， 经过 Flink 做业去支持模型训练和模型的预测。