分散式Tensorflow原理初探

08-16

分散式Tensorflow原理初探

來自專欄銳眼出版社1 人贊了文章

因工作緣故最近正在對深度學習框架Tensorflow的分散式原理髮起猛攻，官網、白皮書和各博客所述稍有偏差，易對新人造成困惑，經過多天推敲，現將心得總結如下，以記錄想法為目的，也望大神不吝賜教。

分散式TF採用PS架構，集群中的節點分為parameter server節點和worker節點。ps和worker代表的是節點所做的task（詳見下文），每個循環中ps的任務是分發參數給worker，worker完成計算密集型任務後將梯度回傳給ps，ps集中處理各梯度並以此更新參數。

分散式訓練策略有模型並行與數據並行之分，TF中一般採用數據並行，即在各worker節點用相同的流圖計算不同的數據。具體到分散式TF的訓練策略又有圖內複製（In-graph）和圖間複製（Between-graph）兩種複製策略，每種策略又分為同步與非同步兩種方式。

In-graph模式只存在一個Client進程（Client詳見下文介紹），這個Client只創建一個Session，這個Session只構建一個圖，並將計算密集型的部分複製到各個worker節點。而Between-graph模式中每個承擔worker任務的節點中都存在一個Client進程（通常與worker在同一進程中），每個Client都構建一張圖，包含所有參數部分與計算部分。這是二者的主要區別。

同步與非同步比較好理解，同步是指每次參數更新需要等待所有worker上傳梯度再一起更新，非同步是指有worker上傳就更新參數。同步方式比較穩定但非同步方式效率較高且有一定的跳出局部最小的作用。因分散式集群主要想解決的也是訓練時長問題，故目前TF中常採用的是非同步Between-graph.

目前一些工程師分享的博客中常用下圖來介紹分散式TF與單機TF的區別，此處極易造成迷惑。此圖源於TF白皮書2015，當時並無In-graph與Between-graph之分，故並未說明右圖實為In-graph模式，因為Between-graph模式中各Worker Service中都應有獨立的Master Service和Client.

TF集群由Server、Client組成，各組成部分說明如下：

Server是伺服器進程，一個運行tf.train.Server的實例，是分散式TF集群的重要組成部分，TF Server採用主從模式，一個Server可輸出一個Master Service和一個Worker Service。

Master

Service是一個RPC（遠程過程調用）服務，作用有三個：一是實現集群設備的遠程訪問；二是作為session的target，它也是tf::Session的介面，三是負責一或多個Worker Service的任務分配。一個Master對對應多個Worker的情況是這樣的，在In-graph模式下，比如一機多卡，每個GPU卡都可以啟動一個Server並輸出Master

Service和Worker Service，但只有一個Server輸出的Master Service起實質作用，作為Session target並負責與其他Worker溝通工作。

Worker

Service也是一個RPC服務，調配本地設備實現子圖的計算。每一個圖計算任務稱為task，task又分為ps（參數伺服器）task和worker task（計算密集型任務處理），ps和worker的集合各自組成一個job。

Client構成TF的計算圖並構建tf::Session來與集群會話。在分散式集群中，創建session需要指定target（target默認值在本地故單機模式無需指明），由上文所提的Master Service提供（通過server.target調用）。一個Client可與多個Server聯接（In-graph模式）；一個Server也可以服務多個Client，也就是在Client構建session時將target同時指定一個server.

Server與task是一一對應的，在構建節點時時使用tf.train.ClusterSpec方法，通過ip:port定義server所在位置，並指定task類型（所屬job集合）。

cluster = tf.train.ClusterSpec({"worker": ["worker0.example.com:2222","worker1.example.com:2222","worker2.example.com:2222"],"ps": ["ps0.example.com:2222","ps1.example.com:2222"]})

之後使用tf.train.Server創建Server，指定集群，jobname和task_index以作為Server節點的身份識別標誌。

server = tf.train.Server(cluster, job_name="worker", task_index=0)

如果是In-graph模式，只需指定一個ip:port來說明在哪個Server的target上創建session，而在Between-graph模式下參數伺服器只需要通過join( )方法使進程進入待命狀態，而worker task節點需要使用tf.train.replica_device_setter( )方法來創建session，並建立ps task與worker task之間的映射，此方法在Between-graph模式中最為關鍵，需要提供cluster，worker_device, ps_device等參數。

with tf.device(tf.train.replica_device_setter( worker_device="/job:worker/task:%d" %FLAGS.task_index,cluster=cluster)):

在replica_device_setter方法中會根據ps_strategy設定參數在ps集群的分配方式，而且會判斷其下的操作是否屬於ps op，從而判斷需要放在負責ps task還是負責worker task的節點上，TF源碼中有具體的邏輯代碼。下圖即為ps op，需要放在負責ps task的節點，所有其他操作均為worker op，放在負責worker task的節點.

Ref:

[1]

http://download.tensorflow.org/paper/whitepaper2015.pdf

[2]

https://tensorflow.google.cn/deploy/distributed