本文翻譯自： 《Multi-GPU processing with data parallelism》， 如有侵權請聯系刪除，僅限于學術交流，請勿商用。如有謬誤，請聯系指出。

如果你使用類似C++這樣的語言在單核CPU上編寫你的軟件，為使其能夠在多個GPU上并行運行，你可能需要從頭開始重寫你的軟件。但是在TensorFlow中并非如此。由于其符號性質，tensorflow可以隱藏所有這些復雜的過程，使你無需在多個CPU和GPU上擴展程序。

讓我們從在CPU上添加兩個向量開始：

import tensorflow as tfwith tf.device(tf.DeviceSpec(device_type="CPU", device_index=0)):a = tf.random_uniform([1000, 100])b = tf.random_uniform([1000, 100])c = a + btf.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 + bc = make_parallel(model, 2, a=a, b=b)復制代碼

你可以使用任何一個將張量作為輸入并返回張量的函數來替換模型，限定條件是輸入和輸出都必須在一個批次（batch）內。值得注意的是，我們還添加了一個變量作用域并將reuse屬性設置為true。這個操作確保我們可以使用相同的變量來處理兩個部分的數據。如此操作讓我們在下一個例子中變得很方便。

讓我們看一個稍微更實際的例子。我們想在多個GPU上訓練神經網絡。在訓練期間，我們不僅需要計算前向傳播，還需要計算后向傳播（梯度變化）。但是我們如何并行化梯度計算呢？事實證明這很簡單。

回憶一下第一項我們想要把一個二階多項式擬合到一組樣本中。我們對代碼進行了一些重組，以便在模型函數中進行大量的操作:

import numpy as np import tensorflow as tfdef 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 lossx = 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) + 3return x_val, y_valsess = 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。這確保了梯度操作可以在與初始操作相同的設備上運行。

總結

以上是生活随笔為你收集整理的【译】Effective TensorFlow Chapter10——在TensorFlow中利用多GPU处理并行数据的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。