Tensorflow 是一個(gè)為數(shù)值計(jì)算(最常見(jiàn)的是訓(xùn)練神經(jīng)網(wǎng)絡(luò))設(shè)計(jì)的流行開源庫(kù)。在這個(gè)框架中，計(jì)算流程通過(guò)數(shù)據(jù)流程圖(data flow graph)設(shè)計(jì)，這為更改操作結(jié)構(gòu)與安置提供了很大靈活性。TensorFlow 允許多個(gè) worker 并行計(jì)算，這對(duì)必須通過(guò)處理的大量訓(xùn)練數(shù)據(jù)訓(xùn)練的神經(jīng)網(wǎng)絡(luò)是有益的。此外，如果模型足夠大，這種并行化有時(shí)可能是必須的。在本文中，我們將探討 TensorFlow 的分布式計(jì)算機(jī)制。

TensorFlow 計(jì)算圖示例

數(shù)據(jù)并行 VS. 模型并行

當(dāng)在多個(gè)計(jì)算節(jié)點(diǎn)間分配神經(jīng)網(wǎng)絡(luò)訓(xùn)練時(shí)，通常采用兩種策略：數(shù)據(jù)并行和模型并行。在前者中，在每個(gè)節(jié)點(diǎn)上單***建模型的實(shí)例，并饋送不同的訓(xùn)練樣本;這種架構(gòu)允許更高的訓(xùn)練吞吐量。相反，在模型并行中，模型的單一實(shí)例在多個(gè)節(jié)點(diǎn)間分配，這種架構(gòu)允許訓(xùn)練更大的模型(可能不一定適合單節(jié)點(diǎn)的存儲(chǔ)器)。如果需要，也可以組合這兩種策略，使給定模型擁有多個(gè)實(shí)例，每個(gè)實(shí)例跨越多個(gè)節(jié)點(diǎn)。在本文中，我們將重點(diǎn)關(guān)注數(shù)據(jù)并行。

數(shù)據(jù)并行與模型并行的不同形式。左：數(shù)據(jù)并行;中：模型并行;右：數(shù)據(jù)并行與模型并行

TensorFlow 中的數(shù)據(jù)并行

當(dāng)使用 TensorFlow 時(shí)，數(shù)據(jù)并行主要表現(xiàn)為兩種形式：圖內(nèi)復(fù)制(in-graph replication)和圖間復(fù)制(between-graph replication)。兩種策略之間最顯著的區(qū)別在于流程圖的結(jié)構(gòu)與其結(jié)果。

1. 圖內(nèi)復(fù)制

圖內(nèi)復(fù)制通常被認(rèn)為是兩種方法中更簡(jiǎn)單和更直接(但更不可擴(kuò)展的)的方法。當(dāng)采用這種策略時(shí)，需要在分布式的主機(jī)上創(chuàng)建一個(gè)包含所有 worker 設(shè)備中副本的流程圖?？梢韵胂?，隨著 worker 數(shù)量的增長(zhǎng)，這樣的流程圖可能會(huì)大幅擴(kuò)展，這可能會(huì)對(duì)模型性能產(chǎn)生不利影響。然而，對(duì)于小系統(tǒng)(例如，雙 GPU 臺(tái)式計(jì)算機(jī))，由于其簡(jiǎn)單性，圖內(nèi)復(fù)制可能是***的。

以下是使用單個(gè) GPU 的基線 TensorFlow 方法與應(yīng)用圖內(nèi)復(fù)制方法的代碼片段的對(duì)比?？紤]到圖內(nèi)復(fù)制方法與擴(kuò)展(scaling)相關(guān)的問(wèn)題，我們將僅考慮單機(jī)、多 GPU 配置的情況。這兩個(gè)代碼片段之間的差異非常小，它們的差異僅存在于：對(duì)輸入數(shù)據(jù)的分塊，使得數(shù)據(jù)在各 worker 間均勻分配，遍歷每個(gè)含有 worker 流程圖的設(shè)備，并將來(lái)自不同 worker 的結(jié)果連接起來(lái)。通過(guò)少量代碼更改，我們可以利用多個(gè)設(shè)備，這種方法使可擴(kuò)展性不再成為大障礙，從而在簡(jiǎn)單配置下更受歡迎。

# single GPU (baseline)   單個(gè) 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   圖內(nèi)復(fù)制 
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) 

這些更改也可以通過(guò)檢查下面的 TensorFlow 流程圖來(lái)可視化。增加的 GPU 模塊說(shuō)明了原始方法的擴(kuò)展方式。

圖內(nèi)復(fù)制的可視化。左：原始圖。右：圖內(nèi)復(fù)制的結(jié)果圖。

2. 圖間復(fù)制

認(rèn)識(shí)到圖內(nèi)復(fù)制在擴(kuò)展上的局限性，圖間復(fù)制的優(yōu)勢(shì)在于運(yùn)用大量節(jié)點(diǎn)時(shí)保證模型性能。這是通過(guò)在每個(gè) worker 上創(chuàng)建計(jì)算圖的副本來(lái)實(shí)現(xiàn)的，并且不需要主機(jī)保存每個(gè) worker 的圖副本。通過(guò)一些 TensorFlow 技巧來(lái)協(xié)調(diào)這些 worker 的圖——如果兩個(gè)單獨(dú)的節(jié)點(diǎn)在同一個(gè) TensorFlow 設(shè)備上分配一個(gè)具有相同名稱的變量，則這些分配將被合并，變量將共享相同的后端存儲(chǔ)，從而這兩個(gè) worker 將合并在一起。

但是，必須確保設(shè)備的正確配置。如果兩個(gè) worker 在不同的設(shè)備上分配變量，則不會(huì)發(fā)生合并。對(duì)此，TensorFlow 提供了 replica_device_setter 函數(shù)。只要每個(gè) worker 以相同的順序創(chuàng)建計(jì)算圖，replica_device_setter 為變量分配提供了確定的方法，確保變量在同一設(shè)備上。這將在下面的代碼中演示。

由于圖間復(fù)制在很大程度上重復(fù)了原始圖，因此多數(shù)相關(guān)的修改實(shí)際上都在集群中節(jié)點(diǎn)的配置上。因此，下面的代碼段將只針對(duì)這一點(diǎn)進(jìn)行改動(dòng)。重要的是要注意，這個(gè)腳本通常會(huì)在集群中的每臺(tái)機(jī)器上執(zhí)行，但具體的命令行參數(shù)不同。下面來(lái)逐行研究代碼。

import sysimport tensorflow as tf 
# specify the cluster's architecture 
cluster = tf.train.ClusterSpec({'ps': ['192.168.1.1:1111'],'worker': ['192.168.1.2:1111','192.168.1.3:1111']})# parse command-line to specify machine 
job_type = sys.argv[1]  # job type: "worker" or "ps" 
task_idx = sys.argv[2]  # index job in the worker or ps list# as defined in the ClusterSpec# create TensorFlow Server. This is how the machines communicate. 
server = tf.train.Server(cluster, job_name=job_type, task_index=task_idx)# parameter server is updated by remote clients.# will not proceed beyond this if statement.if job_type == 'ps': 
   server.join()else:# workers onlywith tf.device(tf.train.replica_device_setter( 
                       worker_device='/job:worker/task:'+task_idx, 
                       clustercluster=cluster)):# build your model here as if you only were using a single machinewith tf.Session(server.target):# train your model here 

運(yùn)行分布式 TensorFlow 的***步是使用 tf.train.ClusterSpec 來(lái)指定集群的架構(gòu)。節(jié)點(diǎn)通常分為兩個(gè)角色(或「job」)：含有變量的參數(shù)服務(wù)器(「ps」)和執(zhí)行大量計(jì)算的「worker」。下面提供每個(gè)節(jié)點(diǎn)的 IP 地址和端口。接下來(lái)，腳本必須確定其 job 類型和在網(wǎng)絡(luò)中的索引;這通常是通過(guò)將命令行參數(shù)傳遞給腳本并解析來(lái)實(shí)現(xiàn)的。job_type 指定節(jié)點(diǎn)是運(yùn)行 ps 還是 worker 任務(wù)，而 task_idx 指定節(jié)點(diǎn)在 ps 或 worker 列表中的索引。使用以上變量創(chuàng)建 TensorFlow 服務(wù)器，用于連接各設(shè)備。

接下來(lái)，如果節(jié)點(diǎn)是參數(shù)服務(wù)器，它只連接它們的線程并等待它們終止。雖然似乎沒(méi)有特定的 ps 代碼，但圖元素實(shí)際上是由 worker 推送到 ps 的。

相反，如果設(shè)備是 worker，則使用 replica_device_setter 構(gòu)建我們的模型，以便在前面討論的這些 ps 服務(wù)器上連續(xù)分配參數(shù)。這些副本將在很大程度上與單機(jī)的流程圖相同。***，我們創(chuàng)建一個(gè) tf.Session 并訓(xùn)練我們的模型。

總結(jié)

希望本文清楚地闡述了與分布式 TensorFlow 相關(guān)的一些術(shù)語(yǔ)和技術(shù)。在以后的文章中，我們將詳細(xì)探討與此相關(guān)及其它的主題。

原文：https://clindatsci.com/blog/2017/5/31/distributed-tensorflow

【本文是51CTO專欄機(jī)構(gòu)“機(jī)器之心”的原創(chuàng)譯文，微信公眾號(hào)“機(jī)器之心( id: almosthuman2014)”】

戳這里，看該作者更多好文