透過多 GPU 並行的方式可以有很好的加速效果，然而一臺機器上所支援的 GPU 是有限的，因此本文介紹了分散式 TensorFlow。分散式 TensorFlow 允許我們在多臺機器上執行一個模型，所以訓練速度或加速效果能顯著地提升。本文簡要介紹了分散式 TensorFlow 的原理與實踐，希望能為準備入坑分散式訓練的讀者提供簡要的介紹。
不幸的是，關於分散式 TensorFlow 的官方文件過於簡略。我們需要一個稍微易懂的介紹，即透過 Jupyter 執行一些基本例子。
如果你想互動式地使用 Jupyter，可以在 GitHub 上找到原始碼。此外，本文的一些解釋是作者自己對實證結果或 TensorFlow 文件的解釋，因此可能會有一些小誤差。
GitHub 地址：https://github.com/mrahtz/distributed_tensorflow_a_gentle_introduction
簡介
import tensorflow as tf比方說，我們希望多個程式共享一些共同的引數。為了簡單起見，假設這只是一個單一的變數：
var = tf.Variable(initial_value=0.0)第一步，我們需要為每個程式建立自己的會話。（假設 sess1 在一個程式中建立，而 sess2 會在另一個程式中建立）。
sess1 = tf.Session()
sess2 = tf.Session()
sess1.run(tf.global_variables_initializer())
sess2.run(tf.global_variables_initializer())每次呼叫 tf.Session() 都會建立一個單獨的「執行引擎」，然後將會話控制程式碼連線到執行引擎。執行引擎是實際儲存變數值並執行操作的東西。且 Python 天生是物件導向的程式設計，它裡面的元素都是類或物件，因此更正式地說，tf.Seesio() 是 TensorFlow 中的一個方法，它會開啟一個會話並執行計算圖。
通常，不同程式中的執行引擎是不相關的。在一個會話中更改變數（在一個執行引擎上）不會影響其他會話中的變數。
print("Initial value of var in session 1:", sess1.run(var))
print("Initial value of var in session 2:", sess2.run(var))
sess1.run(var.assign_add(1.0))
print("Incremented var in session 1")
print("Value of var in session 1:", sess1.run(var))
print("Value of var in session 2:", sess2.run(var))上面程式碼塊的輸出結果為：
Initial value of var in session 1: 0.0
Initial value of var in session 2: 0.0
Incremented var in session 1
Value of var in session 1: 1.0
Value of var in session 2: 0.0對於分散式 TensorFlow，我們首先需要了解它的基本原理。以下程式碼展示瞭如何構建一個最簡單 TensorFlow 叢集，以幫助我們理解它的基本原理。
import tensorflow as tf
c = tf.constant("Hello, Distributed TensorFlow!")
# 建立一個本地TensorFlow叢集
server = tf.train.Server.create_local_server()
# 在叢集上建立一個會話
sess = tf.Session(server.target)
print(sess.run(c)）在以上程式碼中，我們先透過 tf.train.Server.create_local_server 在本地建立一個只有一臺機器的 TensorFlow 叢集。然後在叢集上生成一個會話，透過該對話，我們可以將建立的計算圖執行在 TensorFlow 叢集上。雖然這只是一個單機叢集，但它基本上反映了 TensorFlow 叢集的工作流程。
TensorFlow 叢集會透過一系列任務（task）來執行計算圖中的運算，一般來說不同的任務會在不同的機器上執行。TensorFlow 叢集中的任務也會被聚集為工作（job）。例如在訓練深度模型時，一臺執行反向傳播的機器是一個任務，而所有執行反向傳播的集合是一個工作。上面簡單的案例只是一個任務的叢集，若一個 TensorFlow 叢集有多個任務時，我們需要使用 tf.train.ClusterSpec 來指定每一個任務的機器。
使用分散式 TensorFlow 訓練深度學習模型一般有兩種方式，即 in-graph replication 和 between-graph replication。第一種計算圖內的分散式會令所有任務都使用一個 TensorFlow 計算圖中的變數，而只是將計算部分分配到不同的伺服器上。而另一種計算圖間的分散式會在每一個計算伺服器上建立一個獨立的 TensorFlow 計算圖，但不同計算圖中的相同引數需要以一種固定的方式存放到同一個引數伺服器中。以上大概就是分散式 TensorFlow 的基本概念，隨後我們將透過具體的案例與程式碼加深這一部分的理解。
分散式 TensorFlow
為了在程式之間共享變數，我們需要將不同的執行引擎連線在一起，並輸入分散式張量流。
若使用分散式 TensorFlow，每個程式會執行一個特殊的執行引擎：一個 TensorFlow 伺服器。伺服器作為叢集的一部分連結在一起。（群集中的每個伺服器也稱為任務。）
第一步是定義叢集的規模。我們從最簡單的叢集開始：即兩臺伺服器（兩個任務），它們都在同一臺機器上，一個在 2222 埠，一個在 2223 埠。
tasks = ["localhost:2222", "localhost:2223"]每個任務都與「工作」（job）相關聯，該工作是相關任務的集合。我們將這兩個任務與一個稱為「local」的工作相關聯。
jobs = {"local": tasks}所有這些即定義為一個叢集。
cluster = tf.train.ClusterSpec(jobs)我們現在可以啟動伺服器，指定每個伺服器對應為叢集定義中的哪個伺服器。立即啟動各伺服器，監聽叢集設定中指定的埠。
# "This server corresponds to the the first task (task_index=0)
# of the tasks associated with the 'local' job."
server1 = tf.train.Server(cluster, job_name="local", task_index=0)
server2 = tf.train.Server(cluster, job_name="local", task_index=1)將伺服器連線在同一個叢集中，我們現在可以體驗到分散式 TensorFlow 的強大功能：任何具有相同名稱的變數都將在所有伺服器之間共享。
最簡單的例子是在所有的伺服器上執行同一張靜態計算圖，且每個圖只有一個變數：
tf.reset_default_graph()
var = tf.Variable(initial_value=0.0, name='var')
sess1 = tf.Session(server1.target)
sess2 = tf.Session(server2.target)現在，在一臺伺服器上對變數所作的修改將在第二臺伺服器上作映象處理。
sess1.run(tf.global_variables_initializer())
sess2.run(tf.global_variables_initializer())
print("Initial value of var in session 1:", sess1.run(var))
print("Initial value of var in session 2:", sess2.run(var))
sess1.run(var.assign_add(1.0))
print("Incremented var in session 1")
print("Value of var in session 1:", sess1.run(var))
print("Value of var in session 2:", sess2.run(var))
Initial value of var in session 1: 0.0
Initial value of var in session 2: 0.0
Incremented var in session 1
Value of var in session 1: 1.0
Value of var in session 2: 1.0請注意，因為我們只有一個變數且該變數由兩個會話共享，第二個會話再呼叫 global_variables_initializer 就有些多餘。
存放
現在我們可能會想：變數究竟儲存在哪個伺服器上？又是哪個伺服器在執行操作？
按經驗來說，變數和操作都預設儲存在叢集的第一個任務上。
def run_with_location_trace(sess, op):
# From https://stackoverflow.com/a/41525764/7832197
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
sess.run(op, options=run_options, run_metadata=run_metadata)
for device in run_metadata.step_stats.dev_stats:
print(device.device)
for node in device.node_stats:
print(" ", node.node_name)例如，如果我們使用連線到第一個任務的會話來處理變數 var，那麼所有操作都會執行在這個任務上：
run_with_location_trace(sess1, var)
/job:local/replica:0/task:0/device:CPU:0
_SOURCE
var
run_with_location_trace(sess1, var.assign_add(1.0))
/job:local/replica:0/task:0/device:CPU:0
_SOURCE
AssignAdd_1/value
var
AssignAdd_1但是，如果我們嘗試使用連線到第二個任務的會話處理變數 var，那麼圖節點仍然會在第一個任務上執行。
run_with_location_trace(sess2, var)
/job:local/replica:0/task:1/device:CPU:0
_SOURCE
/job:local/replica:0/task:0/device:CPU:0
_SOURCE
var要將一個變數或操作固定到特定任務上，我們可以使用 tf.device：
with tf.device("/job:local/task:0"):
var1 = tf.Variable(0.0, name='var1')
with tf.device("/job:local/task:1"):
var2 = tf.Variable(0.0, name='var2')
# (This will initialize both variables)
sess1.run(tf.global_variables_initializer())
現在，var1 像之前一樣執行在第一個任務上。
run_with_location_trace(sess1, var1)
/job:local/replica:0/task:0/device:CPU:0
_SOURCE
var1但是 var2 執行在第二個任務上。即使我們嘗試使用連線到第一個任務的會話來評估它，它仍然在第二個任務上執行。
run_with_location_trace(sess1, var2)
/job:local/replica:0/task:0/device:CPU:0
_SOURCE
/job:local/replica:0/task:1/device:CPU:0
_SOURCE
var2變數 2 亦是如此。
run_with_location_trace(sess2, var2)
/job:local/replica:0/task:1/device:CPU:0
_SOURCE
var2
run_with_location_trace(sess2, var1)
/job:local/replica:0/task:1/device:CPU:0
_SOURCE
/job:local/replica:0/task:0/device:CPU:0
_SOURCE
var1計算圖
分散式 TensorFlow 處理圖的過程有幾點需要注意。
誰構建了這個圖？首先，儘管在整個叢集中共享變數值，但圖並不會自動共享。我們用兩臺伺服器建立一個新的叢集，然後用顯式建立的圖設定第一臺伺服器。
cluster = tf.train.ClusterSpec({"local": ["localhost:2224", "localhost:2225"]})
server1 = tf.train.Server(cluster, job_name="local", task_index=0)
server2 = tf.train.Server(cluster, job_name="local", task_index=1)
graph1 = tf.Graph()
with graph1.as_default():
var1 = tf.Variable(0.0, name='var')
sess1 = tf.Session(target=server1.target, graph=graph1)
print(graph1.get_operations())
[, , , ]如果我們建立連線到第二臺伺服器的會話，請注意圖不會自動獲取映象。
graph2 = tf.Graph()
sess2 = tf.Session(target=server2.target, graph=graph2)
print(graph2.get_operations())
————————————————————————————
[]要訪問共享變數，我們必須手動新增一個同名的變數到第二個圖中。
with graph2.as_default():
var2 = tf.Variable(0.0, name='var')
只有如此我們才可以訪問它。
sess1.run(var1.assign(1.0))
sess2.run(var2)
————————————————————————————
1.0關鍵是：每個伺服器負責建立自己的圖。
所有伺服器上的圖都必須一樣嗎？
到目前為止，我們所有的例子都是在兩臺伺服器上執行相同的圖。這被稱為圖內複製（in-graph replication）。
例如，假設我們有一個包含三臺伺服器的叢集。伺服器 1 儲存共享引數，而伺服器 2 和伺服器 3 是工作站節點，每個都有本地變數。在圖內複製中，每臺伺服器的圖如下所示：
圖內複製的問題在於每個伺服器都必須具有整個圖的副本，包括可能只與其他伺服器相關的子圖。這可能會導致圖變得非常大。
另一種方法是圖間複製（between-graph replication）。在這裡，每個伺服器都執行一個只包含共享引數的圖，而且任何變數和操作都與單個伺服器相關。
這種方法縮減了圖的大小，因此我們推薦使用圖間複製。
實踐細節
在介紹完整示例之前，有幾個實踐中遇到的細節問題需要討論一下。
如果在所有伺服器互聯之前嘗試在叢集上執行某些程式，會發生什麼？我們再次建立一個雙任務叢集。
cluster = tf.train.ClusterSpec({
"local": ["localhost:2226", "localhost:2227"]
})這一次，讓我們在隔離程式中啟動每個伺服器。（這允許我們隨時關閉伺服器，以便再次啟動它們進行後續的實驗。除了關閉啟動伺服器的程式之外，目前沒有其它辦法關閉伺服器。）
from multiprocessing import Process
from time import sleep
def s1():
server1 = tf.train.Server(cluster,
job_name="local",
task_index=0)
sess1 = tf.Session(server1.target)
print("server 1: running no-op...")
sess1.run(tf.no_op())
print("server 1: no-op run!")
server1.join() # Block
def s2():
for i in range(3):
print("server 2: %d seconds left before connecting..."
% (3 - i))
sleep(1.0)
server2 = tf.train.Server(cluster,
job_name="local",
task_index=1)
print("server 2: connected!")
server2.join() # Block
# daemon=True so that these processes will definitely be killed
# when the parent process restarts
p1 = Process(target=s1, daemon=True)
p2 = Process(target=s2, daemon=True)伺服器 1 即刻加入叢集，但伺服器 2 在連線之前等待了一會兒。結果如下所示：
p1.start()
p2.start()
server 2: 3 seconds left before connecting...
server 1: running no-op...
server 2: 2 seconds left before connecting...
server 2: 1 seconds left before connecting...
server 2: connected!
server 1: no-op run!可以看出，每個伺服器都試圖在叢集上執行一個操作，直到所有的伺服器都加入。
p1.terminate()
p2.terminate()當伺服器脫離叢集會怎樣？
我們用兩臺伺服器建立一個叢集。伺服器 1 只是反覆嘗試和執行位於伺服器 1 上的 no-op 操作。伺服器 2 將在兩秒鐘後當機。
def s1():
server1 = tf.train.Server(cluster,
job_name="local",
task_index=0)
with tf.device("/job:local/task:0"):
no_op = tf.no_op()
sess1 = tf.Session(server1.target)
for _ in range(6):
print("Server 1: about to run no-op...", end="")
sess1.run(no_op)
print("success!")
sleep(1.0)
def s2():
server2 = tf.train.Server(cluster,
job_name="local",
task_index=1)
sleep(2.0)
print("Server 2 dieing...")
p1 = Process(target=s1, daemon=True)
p2 = Process(target=s2, daemon=True)
p1.start()
p2.start()
————————————————————————————————
Server 1: about to run no-op...success!
Server 1: about to run no-op...success!
Server 2 dieing...
Server 1: about to run no-op...success!
Server 1: about to run no-op...success!
Server 1: about to run no-op...success!
Server 1: about to run no-op...success!短期內，只要我們試圖執行的操作不在脫離的伺服器上，似乎不會出現問題（我沒有測試過長期執行會發生什麼）。如果操作是在脫離的伺服器上：
def s1():
server1 = tf.train.Server(cluster,
job_name="local",
task_index=0)
# This time, we place the no-op on server 2,
# which is going to leave
with tf.device("/job:local/task:1"):
no_op = tf.no_op()
sess1 = tf.Session(server1.target)
for _ in range(5):
print("Server 1: about to run no-op...", end="")
sess1.run(no_op)
print("success!")
sleep(1.0)
p1 = Process(target=s1, daemon=True)
p2 = Process(target=s2, daemon=True)
p1.start()
p2.start()
——————————————————————————————————
—
Server 1: about to run no-op...success!
Server 1: about to run no-op...success!
Server 2 dieing...然後嘗試執行操作程式碼。
p1.terminate()
p2.terminate()如果伺服器又加入叢集會怎樣？
p1 = Process(target=s1, daemon=True)
p2 = Process(target=s2, daemon=True)
p1.start()
p2.start()
sleep(3.0)
# At this point, server 1 is blocked, and server 2 is dead.
print("Restarting server 2...")
p2 = Process(target=s2, daemon=True)
p2.start()
————————————————————————————
Server 1: about to run no-op...success!
Server 1: about to run no-op...success!
Server 2 dieing...
Restarting server 2...
Process Process-7:
Traceback (most recent call last):
File "/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1323, in _do_call
return fn(*args)
File "/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1302, in _run_fn
status, run_metadata)
File "/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/errors_impl.py", line 473, in __exit__
c_api.TF_GetCode(self.status.status))
tensorflow.python.framework.errors_impl.AbortedError: Graph handle is not found: 0000000000000001
Server 1: about to run no-op...Server 2 dieing...系統報了一個 Graph handle is not found 的錯誤。
因此分散式 TensorFlow 不會自動恢復伺服器故障。（如果您對容錯有興趣，請檢視 https://www.youtube.com/watch?v=la_M6bCV91M。）
誰負責初始化共享變數？
一種方法是讓所有工作站執行 tf.global_variables_initializer()。
但是如果我們想保持程式碼整潔並且只用一個伺服器進行初始化，那麼如果有其他伺服器在初始化之前嘗試使用這些變數，可能會遇到問題。一個解決方案就是讓其他工作站等待，直到使用 tf.report_uninitialized_variables 的初始化開始。
def s1():
server1 = tf.train.Server(cluster,
job_name="local",
task_index=0)
var = tf.Variable(0.0, name='var')
sess1 = tf.Session(server1.target)
print("Server 1: waiting for connection...")
sess1.run(tf.report_uninitialized_variables())
while len(sess1.run(tf.report_uninitialized_variables())) > 0:
print("Server 1: waiting for initialization...")
sleep(1.0)
print("Server 1: variables initialized!")
def s2():
server2 = tf.train.Server(cluster,
job_name="local",
task_index=1)
var = tf.Variable(0.0, name='var')
sess2 = tf.Session(server2.target)
for i in range(3):
print("Server 2: waiting %d seconds before initializing..."
% (3 - i))
sleep(1.0)
sess2.run(tf.global_variables_initializer())
p1 = Process(target=s1, daemon=True)
p2 = Process(target=s2, daemon=True)
p1.start()
p2.start()
—————————————————————————————————
Server 1: waiting for connection...
Server 2: waiting 3 seconds before initializing...
Server 1: waiting for initialization...
Server 2: waiting 2 seconds before initializing...
Server 1: waiting for initialization...
Server 2: waiting 1 seconds before initializing...
Server 1: waiting for initialization...
Server 1: variables initialized!
p1.terminate()
p2.terminate()示例
讓我們把所學的知識融合到最後一個使用多程式的例子中。
我們將建立：
一個儲存單個變數 var 的引數伺服器。兩個工作站任務（worker task），每個工作站將多次增加變數 var 的值。我們將讓引數伺服器多輸出幾次 var 的值，以便檢視其變化。
import tensorflow as tf
from multiprocessing import Process
from time import sleep
cluster = tf.train.ClusterSpec({
"worker": [
"localhost:3333",
"localhost:3334",
],
"ps": [
"localhost:3335"
]
})
def parameter_server():
with tf.device("/job:ps/task:0"):
var = tf.Variable(0.0, name='var')
server = tf.train.Server(cluster,
job_name="ps",
task_index=0)
sess = tf.Session(target=server.target)
print("Parameter server: waiting for cluster connection...")
sess.run(tf.report_uninitialized_variables())
print("Parameter server: cluster ready!")
print("Parameter server: initializing variables...")
sess.run(tf.global_variables_initializer())
print("Parameter server: variables initialized")
for i in range(5):
val = sess.run(var)
print("Parameter server: var has value %.1f" % val)
sleep(1.0)
print("Parameter server: blocking...")
server.join()
def worker(worker_n):
with tf.device("/job:ps/task:0"):
var = tf.Variable(0.0, name='var')
server = tf.train.Server(cluster,
job_name="worker",
task_index=worker_n)
sess = tf.Session(target=server.target)
print("Worker %d: waiting for cluster connection..." % worker_n)
sess.run(tf.report_uninitialized_variables())
print("Worker %d: cluster ready!" % worker_n)
while sess.run(tf.report_uninitialized_variables()):
print("Worker %d: waiting for variable initialization..." % worker_n)
sleep(1.0)
print("Worker %d: variables initialized" % worker_n)
for i in range(5):
print("Worker %d: incrementing var" % worker_n)
sess.run(var.assign_add(1.0))
sleep(1.0)
print("Worker %d: blocking..." % worker_n)
server.join()
ps_proc = Process(target=parameter_server, daemon=True)
w1_proc = Process(target=worker, args=(0, ), daemon=True)
w2_proc = Process(target=worker, args=(1, ), daemon=True)ps_proc.start()
————————————————————————————
Parameter server: waiting for cluster connection...
Parameter server: cluster ready!
Parameter server: initializing variables...
Parameter server: variables initialized
Parameter server: var has value 0.0
Parameter server: var has value 2.0
Parameter server: var has value 4.0
Parameter server: var has value 5.0
Parameter server: var has value 7.0
Parameter server: blocking...w1_proc.start()
————————————————————————————————
Worker 0: waiting for cluster connection...
Worker 0: cluster ready!
Worker 0: waiting for variable initialization...
Worker 0: variables initialized
Worker 0: incrementing var
Worker 0: incrementing var
Worker 0: incrementing var
Worker 0: incrementing var
Worker 0: incrementing var
Worker 0: blocking...
w2_proc.start()
———————————————————————————————
Worker 1: waiting for cluster connection...
Worker 1: cluster ready!
Worker 1: waiting for variable initialization...
Worker 1: variables initialized
Worker 1: incrementing var
Worker 1: incrementing var
Worker 1: incrementing var
Worker 1: incrementing var
Worker 1: incrementing var
Worker 1: blocking...
for proc in [w1_proc, w2_proc, ps_proc]:
proc.terminate()總結
透過本文，我們瞭解了：
- 如何將多個 TensorFlow 執行引擎（執行在不同程式或不同機器上）整合為一個叢集，以便共享變數。
- 如何為變數或操作指定伺服器。
- 圖內複製與圖間複製。
- 在所有伺服器互聯之前或在伺服器脫離叢集之後在叢集上執行操作，會發生什麼。
- 如何等待變數被叢集中的另一個任務初始化。
更多資訊和例項，請查閱官方文件：https://www.tensorflow.org/deploy/distributed。
原文連結：http://amid.fish/distributed-tensorflow-a-gentle-introduction