参与:Jane W、黄小天 Tensorflow 是一个为数值计算(最常见的是训练神经网络)设计的流行开源库。在这个框架中,atv,计算流程通过数据流程图(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) 单个 GPU(基线) 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 gpuwith tf.device('/gpu:0'): output = tf.matmul(input_data, b)# create a session and runwith 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 datafor i in range(num_gpus):with tf.device('/gpu:'+str(i)): outputs.append(tf.matmul(inputs[i], b))# merge the results of the deviceswith tf.device('/cpu:0'): output = tf.concat(outputs, axis=0)# create a session and runwith tf.Session() as sess: sess.run(tf.global_variables_initializer())print sess.run(output) 这些更改也可以通过检查下面的 TensorFlow 流程图来可视化。增加的 GPU 模块说明了原始方法的扩展方式。
图内复制的可视化。左:原始图。右:图内复制的结果图。 图间复制 认识到图内复制在扩展上的局限性,图间复制的优势在于运用大量节点时保证模型性能。这是通过在每个 worker 上创建计算图的副本来实现的,并且不需要主机保存每个 worker 的图副本。通过一些 TensorFlow 技巧来协调这些 worker 的图——如果两个单独的节点在同一个 TensorFlow 设备上分配一个具有相同名称的变量,则这些分配将被合并,变量将共享相同的后端存储,atv,从而这两个 worker 将合并在一起。 (责任编辑:本港台直播) |