【译】Effective TensorFlow Chapter10——在TensorFlow中利用多GPU处理并行数据

本文翻译自： 《Multi-GPU processing with data parallelism》 ， 如有侵权请联系删除，仅限于学术交流，请勿商用。如有谬误，请联系指出。

如果你使用类似C++这样的语言在单核CPU上编写你的软件，为使其能够在多个GPU上并行运行，你可能需要从头开始重写你的软件。但是在TensorFlow中并非如此。由于其符号性质，tensorflow可以隐藏所有这些复杂的过程，使你无需在多个CPU和GPU上扩展程序。

让我们从在CPU上添加两个向量开始：

import tensorflow as tf

with tf.device(tf.DeviceSpec(device_type="CPU", device_index=0)):
   a = tf.random_uniform([1000, 100])
   b = tf.random_uniform([1000, 100])
   c = a + b

tf.Session().run(c)
复制代码

同样的事情在GPU上也可以简单地完成：

with tf.device(tf.DeviceSpec(device_type="GPU", device_index=0)):
    a = tf.random_uniform([1000, 100])
    b = tf.random_uniform([1000, 100])
    c = a + b
复制代码

但是，如果我们有两个GPU并希望同时使用它们呢？为此，我们可以把数据分成两份，并让每个GPU单独处理一个部分：

split_a = tf.split(a, 2)
split_b = tf.split(b, 2)

split_c = []
for i in range(2):
    with tf.device(tf.DeviceSpec(device_type="GPU", device_index=i)):
        split_c.append(split_a[i] + split_b[i])

c = tf.concat(split_c, axis=0)
复制代码

让我们以更一般的形式重写它，以便我们可以用任何其他操作集替换添加：

def make_parallel(fn, num_gpus, **kwargs):
    in_splits = {}
    for k, v in kwargs.items():
        in_splits[k] = tf.split(v, num_gpus)

    out_split = []
    for i in range(num_gpus):
        with tf.device(tf.DeviceSpec(device_type="GPU", device_index=i)):
            with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
                out_split.append(fn(**{k : v[i] for k, v in in_splits.items()}))

    return tf.concat(out_split, axis=0)


def model(a, b):
    return a + b

c = make_parallel(model, 2, a=a, b=b)

复制代码

你可以使用任何一个将张量作为输入并返回张量的函数来替换模型，限定条件是输入和输出都必须在一个批次（batch）内。值得注意的是，我们还添加了一个变量作用域并将 reuse 属性设置为true。这个操作确保我们可以使用相同的变量来处理两个部分的数据。如此操作让我们在下一个例子中变得很方便。

让我们看一个稍微更实际的例子。我们想在多个GPU上训练神经网络。在训练期间，我们不仅需要计算前向传播，还需要计算后向传播（梯度变化）。但是我们如何并行化梯度计算呢？事实证明这很简单。

回忆一下第一项我们想要把一个二阶多项式拟合到一组样本中。我们对代码进行了一些重组，以便在模型函数中进行大量的操作:

import numpy as np
import tensorflow as tf

def model(x, y):
    w = tf.get_variable("w", shape=[3, 1])

    f = tf.stack([tf.square(x), x, tf.ones_like(x)], 1)
    yhat = tf.squeeze(tf.matmul(f, w), 1)

    loss = tf.square(yhat - y)
    return loss

x = tf.placeholder(tf.float32)
y = tf.placeholder(tf.float32)

loss = model(x, y)

train_op = tf.train.AdamOptimizer(0.1).minimize(
    tf.reduce_mean(loss))

def generate_data():
    x_val = np.random.uniform(-10.0, 10.0, size=100)
    y_val = 5 * np.square(x_val) + 3
    return x_val, y_val

sess = tf.Session()
sess.run(tf.global_variables_initializer())
for _ in range(1000):
    x_val, y_val = generate_data()
    _, loss_val = sess.run([train_op, loss], {x: x_val, y: y_val})

_, loss_val = sess.run([train_op, loss], {x: x_val, y: y_val})
print(sess.run(tf.contrib.framework.get_variables_by_name("w")))
复制代码

现在让我们使用我们刚刚编写的 make_parallel 函数来并行化这个操作吧。我们只需要从上面的代码中更改两行代码：

loss = make_parallel(model, 2, x=x, y=y)

train_op = tf.train.AdamOptimizer(0.1).minimize(
    tf.reduce_mean(loss),
    colocate_gradients_with_ops=True)
复制代码

要并行化梯度的反向传播，唯一需要改变的是将 colocate_gradients_with_ops 设置为 true 。这确保了梯度操作可以在与初始操作相同的设备上运行。