Tensorflow分布式原理理解

1. Tensorflow 实现原理

 

实现原理

 

TensorFlow有一个重要组件client,顾名思义,就是客户端,它通过Session的接口与master及多个worker相连。其中每一个worker可以与多个硬件设备(device)相连,比如CPU或GPU,并负责管理这些硬件。而master则负责指导所有worker按流程执行计算图。TensorFlow有单机模式和分布式模式两种实现,其中单机指client、master、worker全部在一台机器上的同一个进程中;分布式的版本允许client、master、worker在不同机器的不同进程中,同时由集群调度系统统一管理各项任务。

 

 

图1.1 TensorFlow单机版本和分布式版本的示例图

 

TensorFlow计算图的运行机制

 

Client

 

Client基于TensorFlow的编程接口,构造计算图。此时,TensorFlow并未执行任何计算。直至建立Session会话,并以Session为桥梁,建立Client与后端运行时的通道,将Protobuf格式的GraphDef发送至Distributed Master。也就是说,当Client对OP结果进行求值时,将触发Distributed Master的计算图的执行过程。如下图所示,Client构建了一个简单计算图。它首先将w与x进行矩阵相乘,再与截距b按位相加,最后更新至s。

 

 

图1.2 简单的TensorFlow计算图

 

Distributed Master

 

在分布式的运行时环境中,Distributed Master根据Session.run的Fetching参数,从计算图中反向遍历,找到所依赖的最小子图。然后Distributed Master负责将该子图再次分裂为多个「子图片段」,以便在不同的进程和设备上运行这些「子图片段」。最后,Distributed Master将这些图片段派发给Work Service。随后Work Service启动「本地子图」的执行过程。Distributed Master将会缓存「子图片段」,以便后续执行过程重复使用这些「子图片段」,避免重复计算。

 

 

图1.3 简单的TensorFlow计算图—开始执行

 

执行图计算

 

如上图所示,Distributed Master开始执行计算子图。在执行之前,Distributed Master会实施一系列优化技术,例如「公共表达式消除」,「常量折叠」等。随后,Distributed Master负责任务集的协同,执行优化后的计算子图。

 

 

图1.4 简单的TensorFlow计算图—子图片段

 

子图片段

 

如上图所示,存在一种合理的「子图片段」划分算法。Distributed Master将模型参数相关的OP进行分组,并放置在PS任务上。其他OP则划分为另外一组,放置在Worker任务上执行。

 

 

图1.5 简单的TensorFlow计算图—插入SEND/RECV节点

 

如上图所示,如果计算图的边被任务节点分割,Distributed Master将负责将该边进行分裂,在两个分布式任务之间插入SEND和RECV节点,实现数据的传递。

 

随后,Distributed Master将「子图片段」派发给相应的任务中执行,在Worker Service成为「本地子图」,它负责执行该子图的上的OP。

 

Worker Service

 

对于每个任务,都将存在相应的Worker Service,它主要负责如下3个方面的职责:

 

处理来自Master的请求;

 

调度OP的Kernel实现,执行本地子图;

 

协同任务之间的数据通信。

 

 

图1.6 简单的TensorFlow计算图—执行本地子图

 

执行本地子图

 

Worker Service派发OP到本地设备,执行Kernel的特定。它将尽最大可能地利用多CPU/GPU的处理能力,并发地执行Kernel实现。

 

另外,TensorFlow根据设备类型,对于设备间的SEND/RECV节点进行特化实现:

 

使用cudaMemcpyAsync的API实现本地CPU与GPU设备的数据传输;

 

对于本地的GPU之间则使用端到端的DMA,避免了跨host CPU昂贵的拷贝过程。

 

对于任务之间的数据传递,TensorFlow支持多协议,主要包括:

 

gRPC over TCP

 

RDMA over Converged Ethernet

 

示例代码

 

import tensorflow as tf

 

# 定义常量

 

hello = tf.constant(‘Hello World!’)

 

sess = tf.Session()

 

“””

 

target: (Optional.) The execution engine to connect to.

 

Defaults to using an in-process engine. See

 

@{$distributed$Distributed TensorFlow}

 

“””

 

sess.run(hello)

 

# Terminal out put info:

 

$ b’Hello World!’

 

以上简短的代码其实就已经运用到了Tensorflow分布式结构,只不过该分布式的client、master、worker都是在本地的同一台机器而已。而多台机器实现分布式只需要指定相应的client、master、worker分布在不同的机器上就可以实现了,下面将详细介绍如何在多台机器上实现Tensorflow分布式。

 

2. Tensorflow 分布式之 In-graph replication (图内拷贝)

 

gRPC

 

上面提到了Tensorflow的单机模式是将client、master、worker都放在同一台机器上,而分布式模式就是将三者分布到多台机器上。这时就需要考虑多台机器之前的通信问题了,在Tensorflow的分布式中常使用的通信协议是gRPC协议,gRPC是Google开发的一个开源的RPC(Remote procedure call)协议。该协议允许运行于一台计算机的程序调用另一台计算机的子程序,而程序员无需额外地为这个交互作用编程。

 

In-graph replication

 

Tensorflow训练模型通常需要一些训练参数,训练参数的分发有两种方式:In-graph replication、Between-graph replication。其中In-graph replication方式的数据分发是在一个节点上,这样的好处是配置简单, 其他多机多GPU的计算节点,只要起个join操作, 暴露一个网络接口,等在那里接受任务就好了。 但是这样的坏处是训练数据的分发在一个节点上,要把训练数据分发到不同的机器上, 严重影响并发训练速度。在大数据训练的情况下, 不推荐使用这种模式。

 

示例代码1 自动节点分配策略—-简单的贪婪策略代价模型估计

 

worker_01

 

import tensorflow as tf

 

# 设置worker 地址与端口

 

worker_01 = “172.17.0.2:2222”

 

worker_02 = “172.17.0.3:2222”

 

worker_hosts = [worker_01, worker_02]

 

# Creates a `ClusterSpec`

 

cluster_spec = tf.train.ClusterSpec({“worker”:

 

worker_hosts})

 

# Creates a new server with ClusterSpec、job_name、

 

task_index server = tf.train.Server(cluster_spec,

 

job_name=”worker”, task_index=1)

 

server.join()

 

# Terminal Output

 

I tensorflow/core/distributed_runtime/rpc/grpc_channel.

 

cc:200] Initialize GrpcChannelCache for job worker ->

 

{0 -> localhost:2222, 1 -> 172.17.0.3:2222}

 

I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.

 

cc:221] Started server with target: grpc://localhost:2222

 

worker_02

 

import tensorflow as tf

 

# 设置worker 地址与端口

 

worker_01 = “172.17.0.2:2222”

 

worker_02 = “172.17.0.3:2222”

 

worker_hosts = [worker_01, worker_02]

 

# Creates a `ClusterSpec`

 

cluster_spec = tf.train.ClusterSpec({“worker”:

 

worker_hosts})

 

# Creates a new server with ClusterSpec、job_name、

 

task_index server = tf.train.Server(cluster_spec,

 

job_name=”worker”, task_index=0)

 

server.join()

 

# Terminal Output

 

I tensorflow/core/distributed_runtime/rpc/grpc_channel.

 

cc:200] Initialize GrpcChannelCache for job worker ->

 

{0 -> 172.17.0.2:2222, 1 -> localhost:2222}

 

I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.

 

cc:221] Started server with target: grpc://localhost:2222

 

client

 

import tensorflow as tf

 

# placeholder

 

X = tf.placeholder(“float”)

 

Y = tf.placeholder(“float”)

 

# Creates weight

 

w = tf.Variable(0.0, name=”weight”)

 

# Creates biases

 

b = tf.Variable(0.0, name=”reminder”)

 

# init Variable

 

init_var = tf.global_variables_initializer()

 

# loss function

 

loss = tf.square(Y – tf.multiply(X, w) – b)

 

# Optimizer

 

op = tf.train.GradientDescentOptimizer(0.01).minimize(cost_op)

 

# input parameter

 

train_X = np.linspace(-1, 1, 101)

 

train_Y = 2 * train_X + np.random.randn(*train_X.shape) * 0.33 + 10

 

with tf.Session(“grpc://172.17.0.2:2222”) as sess:

 

sess.run(init_op)

 

for i in range(10):

 

for (x, y) in zip(train_X, train_Y):

 

sess.run(train_op, feed_dict={X: x, Y: y})

 

print(sess.run(w))

 

print(sess.run(b))

 

从上面的示例代码中我们可以看出:这里使用了三台机器在训练,一个clint,两个worker。这里可能会有一些问题:

 

问题1. master在哪里?

 

问题2. 每一个节点的任务怎幺分配?

 

对于问题1,在In-graph replication方式中master其实就是with tf.Session(“grpc://172.17.0.2:2222”) as sess: 这句话中指定的target,也就是说172.17.0.2:2222这台机器就是master,在运行的时候会在这台机器上打印如下日志:

 

I tensorflow/core/distributed_runtime/master_session.cc:1012] Start master session 254ffd62801d1bee with config:

 

对于问题2,

 

把计算已经从单机多GPU,已经扩展到了多机多GPU了, 不过数据分发还是在一个节点。 这些计算节点暴露出来的网络接口,使用起来就跟本机的一个GPU的使用一样, 只要在操作的时候指定tf.device(“/job:worker/task:n”), 就可以向指定GPU一样,把操作指定到一个计算节点上计算,使用起来和多GPU的类似。

 

示例代码2 用户限制的节点分配策略

 

3. Tensorflow 分布式之 Between-graph replication (图间拷贝)

 

between-graph模式下,训练的参数保存在参数服务器, 数据不用分发, 数据分片的保存在各个计算节点, 各个计算节点自己算自己的, 算完了之后, 把要更新的参数告诉参数服务器,参数服务器更新参数。这种模式的优点是不用训练数据的分发了, 尤其是在数据量在TB级的时候, 节省了大量的时间,所以大数据深度学习还是推荐使用between-graph模式。

发表评论

电子邮件地址不会被公开。 必填项已用*标注