【譯】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 。這確保了梯度操作可以在與初始操作相同的裝置上執行。