TensorFlow分布式计算机制解读：以数据并行为重

作者：Neil Tenenholtz

Tensorflow 是一个为数值计算(最常见的是训练神经网络)设计的流行开源库。在这个框架中，计算流程通过数据流程图(data flow graph)设计，这为更改操作结构与安置提供了很大灵活性。TensorFlow 允许多个 worker 并行计算，这对必须通过处理的大量训练数据训练的神经网络是有益的。此外，如果模型足够大，这种并行化有时可能是必须的。在本文中，我们将探讨 TensorFlow 的分布式计算机制。

TensorFlow 计算图示例

数据并行 VS. 模型并行

当在多个计算节点间分配神经网络训练时，通常采用两种策略：数据并行和模型并行。在前者中，在每个节点上单独创建模型的实例，并馈送不同的训练样本;这种架构允许更高的训练吞吐量。相反，在模型并行中，模型的单一实例在多个节点间分配，这种架构允许训练更大的模型(可能不一定适合单节点的存储器)。如果需要，也可以组合这两种策略，使给定模型拥有多个实例，每个实例跨越多个节点。在本文中，我们将重点关注数据并行。

数据并行与模型并行的不同形式。左：数据并行;中：模型并行;右：数据并行与模型并行。

TensorFlow 中的数据并行

当使用 TensorFlow 时，数据并行主要表现为两种形式：图内复制(in-graph replication)和图间复制(between-graph replication)。两种策略之间最显著的区别在于流程图的结构与其结果。

图内复制

图内复制通常被认为是两种方法中更简单和更直接(但更不可扩展的)的方法。当采用这种策略时，需要在分布式的主机上创建一个包含所有 worker 设备中副本的流程图。可以想象，随着 worker 数量的增长，这样的流程图可能会大幅扩展，这可能会对模型性能产生不利影响。然而，对于小系统(例如，双 GPU 台式计算机)，由于其简单性，图内复制可能是最优的。

以下是使用单个 GPU 的基线 TensorFlow 方法与应用图内复制方法的代码片段的对比。考虑到图内复制方法与扩展(scaling)相关的问题，我们将仅考虑单机、多 GPU 配置的情况。这两个代码片段之间的差异非常小，它们的差异仅存在于：对输入数据的分块，使得数据在各 worker 间均匀分配，遍历每个含有 worker 流程图的设备，并将来自不同 worker 的结果连接起来。通过少量代码更改，我们可以利用多个设备，这种方法使可扩展性不再成为大障碍，从而在简单配置下更受欢迎。

				
					# single GPU (baseline)
					import
					tensorflow
					as
					tf
					# place the initial data on the cpu
					with
					tf
					.
					device
					(
					'/cpu:0'
					)
					:
					input_data
					=
					tf
					.
					Variable
					(
					[
					[
					1
					.
					,
					2
					.
					,
					3
					.
					]
					,
					[
					4
					.
					,
					5
					.
					,
					6
					.
					]
					,
					[
					7
					.
					,
					8
					.
					,
					9
					.
					]
					,
					[
					10
					.
					,
					11
					.
					,
					12
					.
					]
					]
					)
					b
					=
					tf
					.
					Variable
					(
					[
					[
					1
					.
					]
					,
					[
					1
					.
					]
					,
					[
					2
					.
					]
					]
					)
					# compute the result on the 0th gpu
					with
					tf
					.
					device
					(
					'/gpu:0'
					)
					:
					output
					=
					tf
					.
					matmul
					(
					input_data
					,
					b
					)
					# create a session and run
					with
					tf
					.
					Session
					(
					)
					as
					sess
					:
					sess
					.
					run
					(
					tf
					.
					global_variables_initializer
					(
					)
					)
					print
					sess
					.
					run
					(
					output
					)
				
			

				
					# in-graph replication
					import
					tensorflow
					as
					tf
					num_gpus
					=
					2
					# place the initial data on the cpu
					with
					tf
					.
					device
					(
					'/cpu:0'
					)
					:
					input_data
					=
					tf
					.
					Variable
					(
					[
					[
					1
					.
					,
					2
					.
					,
					3
					.
					]
					,
					[
					4
					.
					,
					5
					.
					,
					6
					.
					]
					,
					[
					7
					.
					,
					8
					.
					,
					9
					.
					]
					,
					[
					10
					.
					,
					11
					.
					,
					12
					.
					]
					]
					)
					b
					=
					tf
					.
					Variable
					(
					[
					[
					1
					.
					]
					,
					[
					1
					.
					]
					,
					[
					2
					.
					]
					]
					)
					# split the data into chunks for each gpu
					inputs
					=
					tf
					.
					split
					(
					input_data
					,
					num_gpus
					)
					outputs
					=
					[
					]
					# loop over available gpus and pass input data
					for
					i
					in
					range
					(
					num_gpus
					)
					:
					with
					tf
					.
					device
					(
					'/gpu:'
					+
					str
					(
					i
					)
					)
					:
					outputs
					.
					append
					(
					tf
					.
					matmul
					(
					inputs
					[
					i
					]
					,
					b
					)
					)
					# merge the results of the devices
					with
					tf
					.
					device
					(
					'/cpu:0'
					)
					:
					output
					=
					tf
					.
					concat
					(
					outputs
					,
					axis
					=
					0
					)
					# create a session and run
					with
					tf
					.
					Session
					(
					)
					as
					sess
					:
					sess
					.
					run
					(
					tf
					.
					global_variables_initializer
					(
					)
					)
					print
					sess
					.
					run
					(
					output
					)
				
			

这些更改也可以通过检查下面的 TensorFlow 流程图来可视化。增加的 GPU 模块说明了原始方法的扩展方式。

图内复制的可视化。左：原始图。右：图内复制的结果图。

图间复制

认识到图内复制在扩展上的局限性，图间复制的优势在于运用大量节点时保证模型性能。这是通过在每个 worker 上创建计算图的副本来实现的，并且不需要主机保存每个 worker 的图副本。通过一些 TensorFlow 技巧来协调这些 worker 的图——如果两个单独的节点在同一个 TensorFlow 设备上分配一个具有相同名称的变量，则这些分配将被合并，变量将共享相同的后端存储，从而这两个 worker 将合并在一起。

但是，必须确保设备的正确配置。如果两个 worker 在不同的设备上分配变量，则不会发生合并。对此，TensorFlow 提供了 replica_device_setter 函数。只要每个 worker 以相同的顺序创建计算图，replica_device_setter 为变量分配提供了确定的方法，确保变量在同一设备上。这将在下面的代码中演示。

由于图间复制在很大程度上重复了原始图，因此多数相关的修改实际上都在集群中节点的配置上。因此，下面的代码段将只针对这一点进行改动。重要的是要注意，这个脚本通常会在集群中的每台机器上执行，但具体的命令行参数不同。下面来逐行研究代码。

				
					import
					sys
					import
					tensorflow
					as
					tf
					# specify the cluster's architecture
					cluster
					=
					tf
					.
					train
					.
					ClusterSpec
					(
					{
					'ps'
					:
					[
					'192.168.1.1:1111'
					]
					,
					'worker'
					:
					[
					'192.168.1.2:1111'
					,
					'192.168.1.3:1111'
					]
					}
					)
					# parse command-line to specify machine
					job_type
					=
					sys
					.
					argv
					[
					1
					]
					# job type: "worker" or "ps"
					task_idx
					=
					sys
					.
					argv
					[
					2
					]
					# index job in the worker or ps list
					# as defined in the ClusterSpec
					# create TensorFlow Server. This is how the machines communicate.
					server
					=
					tf
					.
					train
					.
					Server
					(
					cluster
					,
					job_name
					=
					job_type
					,
					task_index
					=
					task_idx
					)
					# parameter server is updated by remote clients.
					# will not proceed beyond this if statement.
					if
					job_type
					==
					'ps'
					:
					server
					.
					join
					(
					)
					else
					:
					# workers only
					with
					tf
					.
					device
					(
					tf
					.
					train
					.
					replica_device_setter
					(
					worker_device
					=
					'/job:worker/task:'
					+
					task_idx
					,
					cluster
					=
					cluster
					)
					)
					:
					# build your model here as if you only were using a single machine
					with
					tf
					.
					Session
					(
					server
					.
					target
					)
					:
					# train your model here
				
			

运行分布式 TensorFlow 的第一步是使用 tf.train.ClusterSpec 来指定集群的架构。节点通常分为两个角色(或「job」)：含有变量的参数服务器(「ps」)和执行大量计算的「worker」。下面提供每个节点的 IP 地址和端口。接下来，脚本必须确定其 job 类型和在网络中的索引;这通常是通过将命令行参数传递给脚本并解析来实现的。job_type 指定节点是运行 ps 还是 worker 任务，而 task_idx 指定节点在 ps 或 worker 列表中的索引。使用以上变量创建 TensorFlow 服务器，用于连接各设备。

接下来，如果节点是参数服务器，它只连接它们的线程并等待它们终止。虽然似乎没有特定的 ps 代码，但图元素实际上是由 worker 推送到 ps 的。

相反，如果设备是 worker，则使用 replica_device_setter 构建我们的模型，以便在前面讨论的这些 ps 服务器上连续分配参数。这些副本将在很大程度上与单机的流程图相同。最后，我们创建一个 tf.Session 并训练我们的模型。

总结

希望本文清楚地阐述了与分布式 TensorFlow 相关的一些术语和技术。在以后的文章中，我们将详细探讨与此相关及其它的主题。

End.

转载请注明来自36大数据（36dsj.com)： 36大数据 » TensorFlow分布式计算机制解读：以数据并行为重