本文翻譯自： 《Multi-GPU processing with data parallelism》， 如有侵權(quán)請(qǐng)聯(lián)系刪除，僅限于學(xué)術(shù)交流，請(qǐng)勿商用。如有謬誤，請(qǐng)聯(lián)系指出。

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

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

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) 復(fù)制代碼

同樣的事情在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 復(fù)制代碼

但是，如果我們有兩個(gè)GPU并希望同時(shí)使用它們呢？為此，我們可以把數(shù)據(jù)分成兩份，并讓每個(gè)GPU單獨(dú)處理一個(gè)部分：

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) 復(fù)制代碼

讓我們以更一般的形式重寫它，以便我們可以用任何其他操作集替換添加：

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)復(fù)制代碼

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

讓我們看一個(gè)稍微更實(shí)際的例子。我們想在多個(gè)GPU上訓(xùn)練神經(jīng)網(wǎng)絡(luò)。在訓(xùn)練期間，我們不僅需要計(jì)算前向傳播，還需要計(jì)算后向傳播（梯度變化）。但是我們?nèi)绾尾⑿谢荻扔?jì)算呢？事實(shí)證明這很簡單。

回憶一下第一項(xiàng)我們想要把一個(gè)二階多項(xiàng)式擬合到一組樣本中。我們對(duì)代碼進(jìn)行了一些重組，以便在模型函數(shù)中進(jìn)行大量的操作:

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"))) 復(fù)制代碼

現(xiàn)在讓我們使用我們剛剛編寫的make_parallel函數(shù)來并行化這個(gè)操作吧。我們只需要從上面的代碼中更改兩行代碼：

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) 復(fù)制代碼

要并行化梯度的反向傳播，唯一需要改變的是將colocate_gradients_with_ops設(shè)置為true。這確保了梯度操作可以在與初始操作相同的設(shè)備上運(yùn)行。

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。