分布式TensorFlow入坑指南

姓名:吴兆阳 ?学号:14020199009

转自机器机器之心

嵌牛导读:通过多 GPU 并行的方式可以有很好的加速效果,然而一台机器上所支撑的 GPU 是有限的,因此本文先容了分布式 TensorFlow。分布式 TensorFlow 允许大家在多台机器上运行一个模型,所以训练速度或加速效果能显著地提升。本文简要概述了分布式 TensorFlow 的原理与实践,希翼能为准备入坑分布式训练的读者提供一些先容。不幸的是,关于分布式 TensorFlow 的官方文档过于简略。大家需要一个稍微易懂的先容,即通过 Jupyter 运行一些基本例子。

嵌牛鼻子:分布式TensorFlow

嵌牛提问:什么是分布式TensorFlow,如何入坑?

嵌牛正文:

概况

importtensorflowastf

比方说,大家希翼多个进程共享一些共同的参数。为了简单起见,假设这只是一个单一的变量:

var = tf.Variable(initial_value=0.0)

第一步,大家需要为每个进程创建自己的会话。(假设 sess1 在一个进程中创建,而 sess2 会在另一个进程中创建)。

sess1 = tf.Session()

sess2 = tf.Session()

sess1.run(tf.global_variables_initializer())

sess2.run(tf.global_variables_initializer())

每次调用 tf.Session() 都会创建一个单独的「实行引擎」,然后将会话句柄连接到实行引擎。实行引擎是实际存储变量值并运行操作的东西。且 Python 天生是面向对象的编程,它里面的元素都是类或对象,因此更正式地说,tf.Seesio() 是 TensorFlow 中的一个方法,它会打开一个会话并运行计算图。

通常,不同进程中的实行引擎是不相关的。在一个会话中更改变量(在一个实行引擎上)不会影响其他会话中的变量。

print("Initial value of var in session 1:", sess1.run(var))

print("Initial value of var in session 2:", sess2.run(var))

sess1.run(var.assign_add(1.0))

print("Incremented var in session 1")

print("Value of var in session 1:", sess1.run(var))

print("Value of var in session 2:", sess2.run(var))

上面代码块的输出结果为:

Initialvalue of varinsession1:0.0

Initialvalue of varinsession2:0.0

Incrementedvarinsession1

Valueof varinsession1:1.0

Valueof varinsession2:0.0

对于分布式 TensorFlow,大家首先需要了解它的基本原理。以下代码展示了如何构建一个最简单 TensorFlow 集群,以帮助大家理解它的基本原理。

importtensorflowastf

c=tf.constant("Hello, Distributed TensorFlow!")

# 创建一个本地TensorFlow集群

server=tf.train.Server.create_local_server()

# 在集群上创建一个会话

sess=tf.Session(server.target)

print(sess.run(c))

在以上代码中,大家先通过 tf.train.Server.create_local_server 在本地创建一个只有一台机器的 TensorFlow 集群。然后在集群上生成一个会话,通过该对话,大家可以将创建的计算图运行在 TensorFlow 集群上。虽然这只是一个单机集群,但它基本上反映了 TensorFlow 集群的工作流程。

TensorFlow 集群会通过一系列任务(task)来实行计算图中的运算,一般来说不同的任务会在不同的机器上运行。TensorFlow 集群中的任务也会被聚集为工作(job)。例如在训练深度模型时,一台运行反向传播的机器是一个任务,而所有运行反向传播的集合是一个工作。上面简单的案例只是一个任务的集群,若一个 TensorFlow 集群有多个任务时,大家需要使用 tf.train.ClusterSpec 来指定每一个任务的机器。

使用分布式 TensorFlow 训练深度学习模型一般有两种方式,即 in-graph replication 和 between-graph replication。第一种计算图内的分布式会令所有任务都使用一个 TensorFlow 计算图中的变量,而只是将计算部分分配到不同的服务器上。而另一种计算图间的分布式会在每一个计算服务器上创建一个独立的 TensorFlow 计算图,但不同计算图中的相同参数需要以一种固定的方式存放到同一个参数服务器中。以上大概就是分布式 TensorFlow 的基本概念,随后大家将通过具体的案例与代码加深这一部分的理解。

分布式 TensorFlow

为了在进程之间共享变量,大家需要将不同的实行引擎连接在一起,并输入分布式张量流。

若使用分布式 TensorFlow,每个进程会运行一个特殊的实行引擎:一个 TensorFlow 服务器。服务器作为集群的一部分链接在一起。(群集中的每个服务器也称为任务。)

第一步是定义集群的规模。大家从最简单的集群开始:即两台服务器(两个任务),它们都在同一台机器上,一个在 2222 端口,一个在 2223 端口。

tasks = ["localhost:2222","localhost:2223"]

每个任务都与「工作」(job)相关联,该工作是相关任务的集合。大家将这两个任务与一个称为「local」的工作相关联。

jobs = {"local": tasks}

所有这些即定义为一个集群。

cluster = tf.train.ClusterSpec(jobs)

大家现在可以启动服务器,指定每个服务器对应为集群定义中的哪个服务器。马上启动各服务器,监听集群设置中指定的端口。

# "This server corresponds to the the first task (task_index=0)

# of the tasks associated with the 'local' job."

server1 = tf.train.Server(cluster, job_name="local", task_index=0)

server2 = tf.train.Server(cluster, job_name="local", task_index=1)

将服务器连接在同一个集群中,大家现在可以体验到分布式 TensorFlow 的强大功能:任何具有相同名称的变量都将在所有服务器之间共享。

最简单的例子是在所有的服务器上运行同一张静态计算图,且每个图只有一个变量:

tf.reset_default_graph()

var = tf.Variable(initial_value=0.0, name='var')

sess1 = tf.Session(server1.target)

sess2 = tf.Session(server2.target)

现在,在一台服务器上对变量所作的修改将在第二台服务器上作镜像处理。

sess1.run(tf.global_variables_initializer())

sess2.run(tf.global_variables_initializer())

print("Initial value of var in session 1:", sess1.run(var))

print("Initial value of var in session 2:", sess2.run(var))

sess1.run(var.assign_add(1.0))

print("Incremented var in session 1")

print("Value of var in session 1:", sess1.run(var))

print("Value of var in session 2:", sess2.run(var))

Initialvalue of varinsession1:0.0

Initialvalue of varinsession2:0.0

Incrementedvarinsession1

Valueof varinsession1:1.0

Valueof varinsession2:1.0

请注意,因为大家只有一个变量且该变量由两个会话共享,第二个会话再调用 global_variables_initializer 就有些多余。

存放

现在大家可能会想:变量究竟存储在哪个服务器上?又是哪个服务器在运行操作?

按经验来说,变量和操作都默认存储在集群的第一个任务上。

defrun_with_location_trace(sess, op):

# From https://stackoverflow.com/a/41525764/7832197

run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)

run_metadata = tf.RunMetadata()

sess.run(op, options=run_options, run_metadata=run_metadata)

fordeviceinrun_metadata.step_stats.dev_stats:

print(device.device)

fornodeindevice.node_stats:

print(" ?", node.node_name)

例如,如果大家使用连接到第一个任务的会话来处理变量 var,那么所有操作都会运行在这个任务上:

run_with_location_trace(sess1, var)

/job:local/replica:0/task:0/device:CPU:0

_SOURCE

var

run_with_location_trace(sess1, var.assign_add(1.0))

/job:local/replica:0/task:0/device:CPU:0

_SOURCE

AssignAdd_1/value

var

AssignAdd_1

但是,如果大家尝试使用连接到第二个任务的会话处理变量 var,那么图节点仍然会在第一个任务上运行。

run_with_location_trace(sess2, var)

/job:local/replica:0/task:1/device:CPU:0

_SOURCE

/job:local/replica:0/task:0/device:CPU:0

_SOURCE

var

要将一个变量或操作固定到特定任务上,大家可以使用 tf.device:

withtf.device("/job:local/task:0"):

var1 = tf.Variable(0.0, name='var1')

withtf.device("/job:local/task:1"):

var2 = tf.Variable(0.0, name='var2')

# (This will initialize both variables)

sess1.run(tf.global_variables_initializer())

现在,var1 像之前一样运行在第一个任务上。

run_with_location_trace(sess1, var1)

/job:local/replica:0/task:0/device:CPU:0

_SOURCE

var1

但是 var2 运行在第二个任务上。即使大家尝试使用连接到第一个任务的会话来评估它,它仍然在第二个任务上运行。

run_with_location_trace(sess1, var2)

/job:local/replica:0/task:0/device:CPU:0

_SOURCE

/job:local/replica:0/task:1/device:CPU:0

_SOURCE

var2

变量 2 亦是如此。

run_with_location_trace(sess2, var2)

/job:local/replica:0/task:1/device:CPU:0

_SOURCE

var2

run_with_location_trace(sess2, var1)

/job:local/replica:0/task:1/device:CPU:0

_SOURCE

/job:local/replica:0/task:0/device:CPU:0

_SOURCE

var1

计算图

分布式 TensorFlow 处理图的过程有几点需要注意。

谁构建了这个图?

首先,尽管在整个集群中共享变量值,但图并不会自动共享。

大家用两台服务器创建一个新的集群,然后用显式创建的图设置第一台服务器。

cluster = tf.train.ClusterSpec({"local": ["localhost:2224","localhost:2225"]})

server1 = tf.train.Server(cluster, job_name="local", task_index=0)

server2 = tf.train.Server(cluster, job_name="local", task_index=1)

graph1 = tf.Graph()

withgraph1.as_default():

var1 = tf.Variable(0.0, name='var')

sess1 = tf.Session(target=server1.target, graph=graph1)

print(graph1.get_operations())

[, , , ]

如果大家创建连接到第二台服务器的会话,请注意图不会自动获取镜像。

graph2 = tf.Graph()

sess2 = tf.Session(target=server2.target, graph=graph2)

print(graph2.get_operations())

————————————————————————————

[]

要访问共享变量,大家必须手动添加一个同名的变量到第二个图中。

withgraph2.as_default():

var2 = tf.Variable(0.0, name='var')

只有如此大家才可以访问它。

sess1.run(var1.assign(1.0))

sess2.run(var2)

————————————————————————————

1.0

关键是:每个服务器负责创建自己的图。

所有服务器上的图都必须一样吗?

到目前为止,大家所有的例子都是在两台服务器上运行相同的图。这被称为图内复制(in-graph replication)。

例如,假设大家有一个包含三台服务器的集群。服务器 1 保存共享参数,而服务器 2 和服务器 3 是工作站节点,每个都有本地变量。在图内复制中,每台服务器的图如下所示:

图内复制的问题在于每个服务器都必须具有整个图的副本,包括可能只与其他服务器相关的子图。这可能会导致图变得非常大。

另一种方法是图间复制(between-graph replication)。在这里,每个服务器都运行一个只包含共享参数的图,而且任何变量和操作都与单个服务器相关。

这种方法缩减了图的大小,因此大家推荐使用图间复制。

实践细节

在先容完整示例之前,有几个实践中遇到的细节问题需要讨论一下。

如果在所有服务器互联之前尝试在集群上运行某些程序,会发生什么?

大家再次创建一个双任务集群。

cluster = tf.train.ClusterSpec({

"local": ["localhost:2226","localhost:2227"]

})

这一次,让大家在隔离进程中启动每个服务器。(这允许大家随时关闭服务器,以便再次启动它们进行后续的实验。除了关闭启动服务器的进程之外,目前没有其它办法关闭服务器。)

frommultiprocessingimportProcess

fromtimeimportsleep

defs1():

server1 = tf.train.Server(cluster,

job_name="local",

task_index=0)

sess1 = tf.Session(server1.target)

print("server 1: running no-op...")

sess1.run(tf.no_op())

print("server 1: no-op run!")

server1.join()# Block

defs2():

foriinrange(3):

print("server 2: %d seconds left before connecting..."

% (3- i))

sleep(1.0)

server2 = tf.train.Server(cluster,

job_name="local",

task_index=1)

print("server 2: connected!")

server2.join()# Block

# daemon=True so that these processes will definitely be killed

# when the parent process restarts

p1 =Process(target=s1, daemon=True)

p2 =Process(target=s2, daemon=True)

服务器 1 即刻加入集群,但服务器 2 在连接之前等待了一会儿。结果如下所示:

p1.start()

p2.start()

server2:3seconds left before connecting...

server1: running no-op...

server2:2seconds left before connecting...

server2:1seconds left before connecting...

server2: connected!

server1: no-op run!

可以看出,每个服务器都试图在集群上运行一个操作,直到所有的服务器都加入。

p1.terminate()

p2.terminate()

当服务器脱离集群会怎样?

大家用两台服务器建立一个集群。服务器 1 只是反复尝试和运行位于服务器 1 上的 no-op 操作。服务器 2 将在两秒钟后宕机。

defs1():

server1 = tf.train.Server(cluster,

job_name="local",

task_index=0)

withtf.device("/job:local/task:0"):

no_op = tf.no_op()

sess1 = tf.Session(server1.target)

for_inrange(6):

print("Server 1: about to run no-op...", end="")

sess1.run(no_op)

print("success!")

sleep(1.0)

defs2():

server2 = tf.train.Server(cluster,

job_name="local",

task_index=1)

sleep(2.0)

print("Server 2 dieing...")

p1 =Process(target=s1, daemon=True)

p2 =Process(target=s2, daemon=True)

p1.start()

p2.start()

————————————————————————————————

Server1: about to run no-op...success!

Server1: about to run no-op...success!

Server2dieing...

Server1: about to run no-op...success!

Server1: about to run no-op...success!

Server1: about to run no-op...success!

Server1: about to run no-op...success!

短期内,只要大家试图运行的操作不在脱离的服务器上,似乎不会出现问题。(我没有测试过长期运行会发生什么。)

如果操作是在脱离的服务器上……

defs1():

server1 = tf.train.Server(cluster,

job_name="local",

task_index=0)

# This time, we place the no-op on server 2,

# which is going to leave

withtf.device("/job:local/task:1"):

no_op = tf.no_op()

sess1 = tf.Session(server1.target)

for_inrange(5):

print("Server 1: about to run no-op...", end="")

sess1.run(no_op)

print("success!")

sleep(1.0)

p1 =Process(target=s1, daemon=True)

p2 =Process(target=s2, daemon=True)

p1.start()

p2.start()

——————————————————————————————————

Server1: about to run no-op...success!

Server1: about to run no-op...success!

Server2dieing...

然后尝试运行操作代码。

p1.terminate()

p2.terminate()

如果服务器又加入集群会怎样?

p1 =Process(target=s1, daemon=True)

p2 =Process(target=s2, daemon=True)

p1.start()

p2.start()

sleep(3.0)

# At this point, server 1 is blocked, and server 2 is dead.

print("Restarting server 2...")

p2 =Process(target=s2, daemon=True)

p2.start()

————————————————————————————

Server1: about to run no-op...success!

Server1: about to run no-op...success!

Server2dieing...

Restartingserver2...

ProcessProcess-7:

Traceback(most recent call last):

File"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py", line1323,in_do_call

returnfn(*args)

File"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py", line1302,in_run_fn

status, run_metadata)

File"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/errors_impl.py", line473,in__exit__

c_api.TF_GetCode(self.status.status))

tensorflow.python.framework.errors_impl.AbortedError:Graphhandleisnotfound:0000000000000001

Server1: about to run no-op...Server2dieing...

系统报了一个 Graph handle is not found 的错误。

因此分布式 TensorFlow 不会自动恢复服务器故障。(如果您对容错有兴趣,请查看 https://www.youtube.com/watch?v=la_M6bCV91M。)

谁负责初始化共享变量?

一种方法是让所有工作站运行 tf.global_variables_initializer()。

但是如果大家想保持代码整洁并且只用一个服务器进行初始化,那么如果有其他服务器在初始化之前尝试使用这些变量,可能会遇到问题。一个解决方案就是让其他工作站等待,直到使用 tf.report_uninitialized_variables 的初始化开始。

defs1():

server1 = tf.train.Server(cluster,

job_name="local",

task_index=0)

var = tf.Variable(0.0, name='var')

sess1 = tf.Session(server1.target)

print("Server 1: waiting for connection...")

sess1.run(tf.report_uninitialized_variables())

whilelen(sess1.run(tf.report_uninitialized_variables())) >0:

print("Server 1: waiting for initialization...")

sleep(1.0)

print("Server 1: variables initialized!")

defs2():

server2 = tf.train.Server(cluster,

job_name="local",

task_index=1)

var = tf.Variable(0.0, name='var')

sess2 = tf.Session(server2.target)

foriinrange(3):

print("Server 2: waiting %d seconds before initializing..."

% (3- i))

sleep(1.0)

sess2.run(tf.global_variables_initializer())

p1 =Process(target=s1, daemon=True)

p2 =Process(target=s2, daemon=True)

p1.start()

p2.start()

—————————————————————————————————

Server1: waitingforconnection...

Server2: waiting3seconds before initializing...

Server1: waitingforinitialization...

Server2: waiting2seconds before initializing...

Server1: waitingforinitialization...

Server2: waiting1seconds before initializing...

Server1: waitingforinitialization...

Server1: variables initialized!

p1.terminate()

p2.terminate()

示例

让大家把所学的常识融合到最后一个使用多进程的例子中。

大家将创建:

一个存储单个变量 var 的参数服务器。

两个工作站任务(worker task),每个工作站将多次增加变量 var 的值。

大家将让参数服务器多输出几次 var 的值,以便查看其变化。

importtensorflowastf

frommultiprocessingimportProcess

fromtimeimportsleep

cluster = tf.train.ClusterSpec({

"worker": [

"localhost:3333",

"localhost:3334",

],

"ps": [

"localhost:3335"

]

})

defparameter_server():

withtf.device("/job:ps/task:0"):

var = tf.Variable(0.0, name='var')

server = tf.train.Server(cluster,

job_name="ps",

task_index=0)

sess = tf.Session(target=server.target)

print("Parameter server: waiting for cluster connection...")

sess.run(tf.report_uninitialized_variables())

print("Parameter server: cluster ready!")

print("Parameter server: initializing variables...")

sess.run(tf.global_variables_initializer())

print("Parameter server: variables initialized")

foriinrange(5):

val = sess.run(var)

print("Parameter server: var has value %.1f"% val)

sleep(1.0)

print("Parameter server: blocking...")

server.join()

defworker(worker_n):

withtf.device("/job:ps/task:0"):

var = tf.Variable(0.0, name='var')

server = tf.train.Server(cluster,

job_name="worker",

task_index=worker_n)

sess = tf.Session(target=server.target)

print("Worker %d: waiting for cluster connection..."% worker_n)

sess.run(tf.report_uninitialized_variables())

print("Worker %d: cluster ready!"% worker_n)

whilesess.run(tf.report_uninitialized_variables()):

print("Worker %d: waiting for variable initialization..."% worker_n)

sleep(1.0)

print("Worker %d: variables initialized"% worker_n)

foriinrange(5):

print("Worker %d: incrementing var"% worker_n)

sess.run(var.assign_add(1.0))

sleep(1.0)

print("Worker %d: blocking..."% worker_n)

server.join()

ps_proc =Process(target=parameter_server, daemon=True)

w1_proc =Process(target=worker, args=(0, ), daemon=True)

w2_proc =Process(target=worker, args=(1, ), daemon=True)

ps_proc.start()

————————————————————————————

Parameterserver: waitingforcluster connection...

Parameterserver: cluster ready!

Parameterserver: initializing variables...

Parameterserver: variables initialized

Parameterserver: var has value0.0

Parameterserver: var has value2.0

Parameterserver: var has value4.0

Parameterserver: var has value5.0

Parameterserver: var has value7.0

Parameterserver: blocking...

w1_proc.start()

————————————————————————————————

Worker0: waitingforcluster connection...

Worker0: cluster ready!

Worker0: waitingforvariable initialization...

Worker0: variables initialized

Worker0: incrementing var

Worker0: incrementing var

Worker0: incrementing var

Worker0: incrementing var

Worker0: incrementing var

Worker0: blocking...

w2_proc.start()

———————————————————————————————

Worker1: waitingforcluster connection...

Worker1: cluster ready!

Worker1: waitingforvariable initialization...

Worker1: variables initialized

Worker1: incrementing var

Worker1: incrementing var

Worker1: incrementing var

Worker1: incrementing var

Worker1: incrementing var

Worker1: blocking...

forprocin[w1_proc, w2_proc, ps_proc]:

proc.terminate()

总结

通过本文,大家了解了:

如何将多个 TensorFlow 实行引擎(运行在不同进程或不同机器上)集成为一个集群,以便共享变量。

如何为变量或操作指定服务器。

图内复制与图间复制。

在所有服务器互联之前或在服务器脱离集群之后在集群上运行操作,会发生什么。

如何等待变量被集群中的另一个任务初始化。

推荐阅读更多精彩内容