Tensorflow 是一個為數值計算(最常見的是訓練神經網路)設計的流行開源庫。在這個框架中,計算流程通過資料流程圖(data flow graph)設計,這為更改操作結構與安置提供了很大靈活性。TensorFlow 允許多個 worker 平行計算,這對必須通過處理的大量訓練資料訓練的神經網路是有益的。此外,如果模型足夠大,這種並行化有時可能是必須的。在本文中,我們將探討 TensorFlow 的分散式計算機制。
TensorFlow 計算圖示例
資料並行 VS. 模型並行
當在多個計算節點間分配神經網路訓練時,通常採用兩種策略:資料並行和模型並行。在前者中,在每個節點上單獨建立模型的例項,並饋送不同的訓練樣本;這種架構允許更高的訓練吞吐量。相反,在模型並行中,模型的單一例項在多個節點間分配,這種架構允許訓練更大的模型(可能不一定適合單節點的儲存器)。如果需要,也可以組合這兩種策略,使給定模型擁有多個例項,每個例項跨越多個節點。在本文中,我們將重點關注資料並行。
資料並行與模型並行的不同形式。左:資料並行;中:模型並行;右:資料並行與模型並行。
TensorFlow 中的資料並行
當使用 TensorFlow 時,資料並行主要表現為兩種形式:圖內複製(in-graph replication)和圖間複製(between-graph replication)。兩種策略之間最顯著的區別在於流程圖的結構與其結果。
圖內複製
圖內複製通常被認為是兩種方法中更簡單和更直接(但更不可擴充套件的)的方法。當採用這種策略時,需要在分散式的主機上建立一個包含所有 worker 裝置中副本的流程圖。可以想象,隨著 worker 數量的增長,這樣的流程圖可能會大幅擴充套件,這可能會對模型效能產生不利影響。然而,對於小系統(例如,雙 GPU 臺式計算機),由於其簡單性,圖內複製可能是最優的。
以下是使用單個 GPU 的基線 TensorFlow 方法與應用圖內複製方法的程式碼片段的對比。考慮到圖內複製方法與擴充套件(scaling)相關的問題,我們將僅考慮單機、多 GPU 配置的情況。這兩個程式碼片段之間的差異非常小,它們的差異僅存在於:對輸入資料的分塊,使得資料在各 worker 間均勻分配,遍歷每個含有 worker 流程圖的裝置,並將來自不同 worker 的結果連線起來。通過少量程式碼更改,我們可以利用多個裝置,這種方法使可擴充套件性不再成為大障礙,從而在簡單配置下更受歡迎。
# single GPU (baseline) 單個 GPU(基線)import tensorflow as tf# place the initial data on the cpuwith 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 圖內複製import tensorflow as tfnum_gpus = 2# place the initial data on the cpuwith 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 gpuinputs = 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)
這些更改也可以通過檢查下面的 TensorFlow 流程圖來視覺化。增加的 GPU 模組說明了原始方法的擴充套件方式。
圖內複製的視覺化。左:原始圖。右:圖內複製的結果圖。
圖間複製
認識到圖內複製在擴充套件上的侷限性,圖間複製的優勢在於運用大量節點時保證模型效能。這是通過在每個 worker 上建立計算圖的副本來實現的,並且不需要主機儲存每個 worker 的圖副本。通過一些 TensorFlow 技巧來協調這些 worker 的圖——如果兩個單獨的節點在同一個 TensorFlow 裝置上分配一個具有相同名稱的變數,則這些分配將被合併,變數將共享相同的後端儲存,從而這兩個 worker 將合併在一起。
但是,必須確保裝置的正確配置。如果兩個 worker 在不同的裝置上分配變數,則不會發生合併。對此,TensorFlow 提供了 replica_device_setter 函式。只要每個 worker 以相同的順序建立計算圖,replica_device_setter 為變數分配提供了確定的方法,確保變數在同一裝置上。這將在下面的程式碼中演示。
由於圖間複製在很大程度上重複了原始圖,因此多數相關的修改實際上都在叢集中節點的配置上。因此,下面的程式碼段將只針對這一點進行改動。重要的是要注意,這個指令碼通常會在叢集中的每臺機器上執行,但具體的命令列引數不同。下面來逐行研究程式碼。
import sysimport tensorflow as tf# specify the cluster's architecturecluster = 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 machinejob_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, cluster=cluster)):# build your model here as if you only were using a single machinewith tf.Session(server.target):# train your model here
執行分散式 TensorFlow 的第一步是使用 tf.train.ClusterSpec 來指定叢集的架構。節點通常分為兩個角色(或「job」):含有變數的引數伺服器(「ps」)和執行大量計算的「worker」。下面提供每個節點的 IP 地址和埠。接下來,指令碼必須確定其 job 型別和在網路中的索引;這通常是通過將命令列引數傳遞給指令碼並解析來實現的。job_type 指定節點是執行 ps 還是 worker 任務,而 task_idx 指定節點在 ps 或 worker 列表中的索引。使用以上變數建立 TensorFlow 伺服器,用於連線各裝置。
接下來,如果節點是引數伺服器,它只連線它們的執行緒並等待它們終止。雖然似乎沒有特定的 ps 程式碼,但圖元素實際上是由 worker 推送到 ps 的。
相反,如果裝置是 worker,則使用 replica_device_setter 構建我們的模型,以便在前面討論的這些 ps 伺服器上連續分配引數。這些副本將在很大程度上與單機的流程圖相同。最後,我們建立一個 tf.Session 並訓練我們的模型。
總結
希望本文清楚地闡述了與分散式 TensorFlow 相關的一些術語和技術。在以後的文章中,我們將詳細探討與此相關及其它的主題。
原文連結:https://clindatsci.com/blog/2017/5/31/distributed-tensorflow