# 十二、設備和服務器上的分布式 TensorFlow
> 譯者:[@空白](https://github.com/yhcheer)
>
> 校對者:[@飛龍](https://github.com/wizardforcel)
在第 11 章,我們討論了幾種可以明顯加速訓練的技術:更好的權重初始化,批量標準化,復雜的優化器等等。 但是,即使采用了所有這些技術,在具有單個 CPU 的單臺機器上訓練大型神經網絡可能需要幾天甚至幾周的時間。
在本章中,我們將看到如何使用 TensorFlow 在多個設備(CPU 和 GPU)上分配計算并將它們并行運行(參見圖 12-1)。 首先,我們會先在一臺機器上的多個設備上分配計算,然后在多臺機器上的多個設備上分配計算。

與其他神經網絡框架相比,TensorFlow 對分布式計算的支持是其主要亮點之一。 它使您可以完全控制如何跨設備和服務器分布(或復制)您的計算圖,并且可以讓您以靈活的方式并行和同步操作,以便您可以在各種并行方法之間進行選擇。
我們來看一些最流行的方法來并行執行和訓練一個神經網絡,這讓我們不再需要等待數周才能完成訓練算法,而最終可能只會等待幾個小時。 這不僅可以節省大量時間,還意味著您可以更輕松地嘗試各種模型,并經常重新訓練模型上的新數據。
還有其他很好的并行化例子,包括當我們在微調模型時可以探索更大的超參數空間,并有效地運行大規模神經網絡。
但我們必須先學會走路才能跑步。 我們先從一臺機器上的幾個 GPU 上并行化簡單圖形開始。
## 一臺機器上多設備
只需添加 GPU 卡到單個機器,您就可以獲得主要的性能提升。 事實上,在很多情況下,這就足夠了。 你根本不需要使用多臺機器。 例如,通常在單臺機器上使用 8 個 GPU,而不是在多臺機器上使用 16 個 GPU(由于多機器設置中的網絡通信帶來的額外延遲),可以同樣快地訓練神經網絡。
在本節中,我們將介紹如何設置您的環境,以便 TensorFlow 可以在一臺機器上使用多個 GPU 卡。 然后,我們將看看如何在可用設備上進行分布操作,并且并行執行它們。
### 安裝
為了在多個 GPU 卡上運行 TensorFlow,首先需要確保 GPU 卡具有 NVidia 計算能力(大于或等于3.0)。 這包括 Nvidia 的 Titan,Titan X,K20 和 K40(如果你擁有另一張卡,你可以在 <https://developer.nvidia.com/cuda-gpus> 查看它的兼容性)。
如果您不擁有任何 GPU 卡,則可以使用具有 GPU 功能的主機服務器,如 Amazon AWS。 在 ?igaAvsec 的[博客文章](https://goo.gl/kbge5b)中,提供了在 Amazon AWS GPU 實例上使用 Python 3.5 設置 TensorFlow 0.9 的詳細說明。將它更新到最新版本的 TensorFlow 應該不會太難。 Google 還發布了一項名為 Cloud Machine Learning 的云服務來運行 TensorFlow 圖表。 2016 年 5 月,他們宣布他們的平臺現在包括配備張量處理器(TPU)的服務器,專門用于機器學習的處理器,比許多 GPU 處理 ML 任務要快得多。 當然,另一種選擇只是購買你自己的 GPU 卡。 Tim Dettmers 寫了一篇很棒的博客文章來幫助你選擇,他會定期更新它。
您必須下載并安裝相應版本的 CUDA 和 cuDNN 庫(如果您使用的是 TensorFlow 1.0.0,則為 CUDA 8.0 和 cuDNN 5.1),并設置一些環境變量,以便 TensorFlow 知道在哪里可以找到 CUDA 和 cuDNN。 詳細的安裝說明可能會相當迅速地更改,因此最好按照 TensorFlow 網站上的說明進行操作。
Nvidia 的 CUDA 允許開發者使用支持 CUDA 的 GPU 進行各種計算(不僅僅是圖形加速)。 Nvidia 的 CUDA 深度神經網絡庫(cuDNN)是針對 DNN 的 GPU 加速原語庫。 它提供了常用 DNN 計算的優化實現,例如激活層,歸一化,前向和后向卷積以及池化(參見第 13 章)。 它是 Nvidia Deep Learning SDK 的一部分(請注意,它需要創建一個 Nvidia 開發者帳戶才能下載它)。 TensorFlow 使用 CUDA 和 cuDNN 來控制 GPU 卡并加速計算(見圖 12-2)。

您可以使用`nvidia-smi`命令來檢查 CUDA 是否已正確安裝。 它列出了可用的 GPU 卡以及每張卡上運行的進程:

最后,您必須安裝支持 GPU 的 TensorFlow。 如果你使用`virtualenv`創建了一個獨立的環境,你首先需要激活它:
<pre><code>
$ cd $ML_PATH
# Your ML working directory (e.g., HOME/ml)
$ source env/bin/activate
</code></pre>
然后安裝合適的支持 GPU 的 TensorFlow 版本:
```
$ pip3 install --upgrade tensorflow-gpu
```
現在您可以打開一個 Python shell 并通過導入 TensorFlow 并創建一個會話來檢查 TensorFlow 是否正確檢測并使用 CUDA 和 cuDNN:
```
>>> import tensorflow as tf
I [...]/dso_loader.cc:108] successfully opened CUDA library libcublas.so locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcudnn.so locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcufft.so locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcuda.so.1 locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcurand.so locally
```
```
>>> sess = tf.Session()
[...]
I [...]/gpu_init.cc:102] Found device 0 with properties:
name: GRID K520
major: 3 minor: 0 memoryClockRate (GHz) 0.797
pciBusID 0000:00:03.0
Total memory: 4.00GiB
Free memory: 3.95GiB
I [...]/gpu_init.cc:126] DMA: 0
I [...]/gpu_init.cc:136] 0: Y
I [...]/gpu_device.cc:839] Creating TensorFlow device
(/gpu:0) -> (device: 0, name: GRID K520, pci bus id: 0000:00:03.0)
```
看起來不錯!TensorFlow 檢測到 CUDA 和 cuDNN 庫,并使用 CUDA 庫來檢測 GPU 卡(在這種情況下是 Nvidia Grid K520 卡)。
## 管理 GPU 內存
默認情況下,TensorFlow 會在您第一次運行圖形時自動獲取所有可用 GPU 中的所有 RAM,因此當第一個程序仍在運行時,您將無法啟動第二個 TensorFlow 程序。 如果你嘗試,你會得到以下錯誤:
```
E [...]/cuda_driver.cc:965] failed to allocate 3.66G (3928915968 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY
```
一種解決方案是在不同的 GPU 卡上運行每個進程。 為此,最簡單的選擇是設置`CUDA_VISIBLE_DEVICES`環境變量,以便每個進程只能看到對應的 GPU 卡。 例如,你可以像這樣啟動兩個程序:
```
$ CUDA_VISIBLE_DEVICES=0,1 python3 program_1.py
# and in another terminal:
$ CUDA_VISIBLE_DEVICES=3,2 python3 program_2.py
```
程序 #1 只會看到 GPU 卡 0 和 1(分別編號為 0 和 1),程序 #2 只會看到 GPU 卡 2 和 3(分別編號為 1 和 0)。 一切都會正常工作(見圖 12-3)。

另一種選擇是告訴 TensorFlow 只抓取一小部分內存。 例如,要使 TensorFlow 只占用每個 GPU 內存的 40%,您必須創建一個`ConfigProto`對象,將其`gpu_options.per_process_gpu_memory_fraction`選項設置為 0.4,并使用以下配置創建`session`:
```py
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config)
```
現在像這樣的兩個程序可以使用相同的 GPU 卡并行運行(但不是三個,因為`3×0.4> 1`)。 見圖 12-4。

如果在兩個程序都運行時運行`nvidia-smi`命令,則應該看到每個進程占用每個卡的總 RAM 大約 40%:

另一種選擇是告訴 TensorFlow 只在需要時才抓取內存。 為此,您必須將`config.gpu_options.allow_growth`設置為`True`。但是,TensorFlow 一旦抓取內存就不會釋放內存(以避免內存碎片),因此您可能會在一段時間后內存不足。 是否使用此選項可能難以確定,因此一般而言,您可能想要堅持之前的某個選項。
好的,現在你已經有了一個支持 GPU 的 TensorFlow 安裝。 讓我們看看如何使用它!
## 設備布置操作
TensorFlow 白皮書介紹了一種友好的動態布置器算法,該算法能夠自動將操作分布到所有可用設備上,并考慮到以前運行圖中所測量的計算時間,估算每次操作的輸入和輸出張量的大小, 每個設備可用的 RAM,傳輸數據進出設備時的通信延遲,來自用戶的提示和約束等等。 不幸的是,這種復雜的算法是谷歌內部的,它并沒有在 TensorFlow 的開源版本中發布。它被排除在外的原因似乎是,由用戶指定的一小部分放置規則實際上比動態放置器放置的更有效。 然而,TensorFlow 團隊正在努力改進它,并且最終可能會被開放。
在此之前,TensorFlow都是簡單的放置,它(如其名稱所示)非常基本。
### 簡單放置
無論何時運行圖形,如果 TensorFlow 需要求值尚未放置在設備上的節點,則它會使用簡單放置器將其放置在未放置的所有其他節點上。 簡單放置尊重以下規則:
- 如果某個節點已經放置在圖形的上一次運行中的某個設備上,則該節點將保留在該設備上。
- 否則,如果用戶將一個節點固定到設備上(下面介紹),則放置器將其放置在該設備上。
- 否則,它默認為 GPU#0,如果沒有 GPU,則默認為 CPU。
正如您所看到的,將操作放在適當的設備上主要取決于您。 如果您不做任何事情,整個圖表將被放置在默認設備上。 要將節點固定到設備上,您必須使用`device()`函數創建一個設備塊。 例如,以下代碼將變量`a`和常量`b`固定在 CPU 上,但乘法節點`c`不固定在任何設備上,因此將放置在默認設備上:
```py
with tf.device("/cpu:0"):
a = tf.Variable(3.0)
b = tf.constant(4.0)
c = a * b
```
其中,`"/cpu:0"`設備合計多 CPU 系統上的所有 CPU。 目前沒有辦法在特定 CPU 上固定節點或僅使用所有 CPU 的子集。
### 記錄放置位置
讓我們檢查一下簡單的放置器是否遵守我們剛剛定義的布局約束條件。 為此,您可以將`log_device_placement`選項設置為`True`;這告訴放置器在放置節點時記錄消息。例如:
```
>>> config = tf.ConfigProto()
>>> config.log_device_placement = True
>>> sess = tf.Session(config=config)
I [...] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GRID K520, pci bus id: 0000:00:03.0)
[...]
>>> x.initializer.run(session=sess)
I [...] a: /job:localhost/replica:0/task:0/cpu:0
I [...] a/read: /job:localhost/replica:0/task:0/cpu:0
I [...] mul: /job:localhost/replica:0/task:0/gpu:0
I [...] a/Assign: /job:localhost/replica:0/task:0/cpu:0
I [...] b: /job:localhost/replica:0/task:0/cpu:0
I [...] a/initial_value: /job:localhost/replica:0/task:0/cpu:0
>>> sess.run(c)
12
```
`Info`中以大寫字母`I`開頭的行是日志消息。 當我們創建一個會話時,TensorFlow 會記錄一條消息,告訴我們它已經找到了一個 GPU 卡(在這個例子中是 Grid K520 卡)。 然后,我們第一次運行圖形(在這種情況下,當初始化變量`a`時),簡單布局器運行,并將每個節點放置在分配給它的設備上。正如預期的那樣,日志消息顯示所有節點都放在`"/cpu:0"`上,除了乘法節點,它以默認設備`"/gpu:0"`結束(您可以先忽略前綴:`/job:localhost/replica:0/task:0;`我們將在一會兒討論它)。 注意,我們第二次運行圖(計算`c`)時,由于 TensorFlow 需要計算的所有節點`c`都已經放置,所以不使用布局器。
### 動態放置功能
創建設備塊時,可以指定一個函數,而不是設備名稱。TensorFlow 會調用這個函數來進行每個需要放置在設備塊中的操作,并且該函數必須返回設備的名稱來固定操作。 例如,以下代碼將固定所有變量節點到`"/cpu:0"`(在本例中只是變量`a`)和所有其他節點到`"/gpu:0"`:
```py
def variables_on_cpu(op):
if op.type == "Variable":
return "/cpu:0"
else:
return "/gpu:0"
with tf.device(variables_on_cpu):
a = tf.Variable(3.0)
b = tf.constant(4.0)
c = a * b
```
您可以輕松實現更復雜的算法,例如以循環方式用GPU鎖定變量。
### 操作和內核
對于在設備上運行的 TensorFlow 操作,它需要具有該設備的實現;這被稱為內核。 許多操作對于 CPU 和 GPU 都有內核,但并非全部都是。 例如,TensorFlow 沒有用于整數變量的 GPU 內核,因此當 TensorFlow 嘗試將變量i放置到 GPU#0 時,以下代碼將失敗:
```
>>> with tf.device("/gpu:0"):
... i = tf.Variable(3)
[...]
>>> sess.run(i.initializer)
Traceback (most recent call last):
[...]
tensorflow.python.framework.errors.InvalidArgumentError: Cannot assign a device to node 'Variable': Could not satisfy explicit device specification
```
請注意,TensorFlow 推斷變量必須是`int32`類型,因為初始化值是一個整數。 如果將初始化值更改為 3.0 而不是 3,或者如果在創建變量時顯式設置`dtype = tf.float32`,則一切正常。
### 軟放置
默認情況下,如果您嘗試在操作沒有內核的設備上固定操作,則當 TensorFlow 嘗試將操作放置在設備上時,您會看到前面顯示的異常。 如果您更喜歡 TensorFlow 回退到 CPU,則可以將`allow_soft_placement`配置選項設置為`True`:
```py
with tf.device("/gpu:0"):
i = tf.Variable(3)
config = tf.ConfigProto()
config.allow_soft_placement = True
sess = tf.Session(config=config)
sess.run(i.initializer) # the placer runs and falls back to /cpu:0
```
到目前為止,我們已經討論了如何在不同設備上放置節點。 現在讓我們看看 TensorFlow 如何并行運行這些節點。
## 并行運行
當 TensorFlow 運行圖時,它首先找出需要求值的節點列表,然后計算每個節點有多少依賴關系。 然后 TensorFlow 開始求值具有零依賴關系的節點(即源節點)。 如果這些節點被放置在不同的設備上,它們顯然會被并行求值。 如果它們放在同一個設備上,它們將在不同的線程中進行求值,因此它們也可以并行運行(在單獨的 GPU 線程或 CPU 內核中)。
TensorFlow 管理每個設備上的線程池以并行化操作(參見圖 12-5)。 這些被稱為 inter-op 線程池。 有些操作具有多線程內核:它們可以使用其他線程池(每個設備一個)稱為 intra-op 線程池(下面寫成內部線程池)。

例如,在圖 12-5 中,操作`A`,`B`和`C`是源操作,因此可以立即進行求值。 操作`A`和`B`放置在 GPU#0 上,因此它們被發送到該設備的內部線程池,并立即進行并行求值。 操作A正好有一個多線程內核; 它的計算被分成三部分,這些部分由內部線程池并行執行。 操作`C`轉到 GPU#1 的內部線程池。
一旦操作`C`完成,操作`D`和`E`的依賴性計數器將遞減并且都將達到 0,因此這兩個操作將被發送到操作內線程池以執行。
您可以通過設置`inter_op_parallelism_threads`選項來控制內部線程池的線程數。 請注意,您開始的第一個會話將創建內部線程池。 除非您將`use_per_session_threads`選項設置為`True`,否則所有其他會話都將重用它們。 您可以通過設置`intra_op_parallelism_threads`選項來控制每個內部線程池的線程數。
## 控制依賴關系
在某些情況下,即使所有依賴的操作都已執行,推遲對操作的求值可能也是明智之舉。例如,如果它使用大量內存,但在圖形中只需要更多內存,則最好在最后一刻對其進行求值,以避免不必要地占用其他操作可能需要的 RAM。 另一個例子是依賴位于設備外部的數據的一組操作。 如果它們全部同時運行,它們可能會使設備的通信帶寬達到飽和,并最終導致所有等待 I/O。 其他需要傳遞數據的操作也將被阻止。 順序執行這些通信繁重的操作將是比較好的,這樣允許設備并行執行其他操作。
推遲對某些節點的求值,一個簡單的解決方案是添加控制依賴關系。 例如,下面的代碼告訴 TensorFlow 僅在求值完`a`和`b`之后才求值`x`和`y`:
```py
a = tf.constant(1.0)
b = a + 2.0
with tf.control_dependencies([a, b]):
x = tf.constant(3.0)
y = tf.constant(4.0)
z = x + y
```
顯然,由于`z`依賴于`x`和`y`,所以求值`z`也意味著等待`a`和`b`進行求值,即使它并未顯式存在于`control_dependencies()`塊中。 此外,由于`b`依賴于`a`,所以我們可以通過在`[b]`而不是`[a,b]`上創建控制依賴關系來簡化前面的代碼,但在某些情況下,“顯式比隱式更好”。
很好!現在你知道了:
- 如何以任何您喜歡的方式在多個設備上進行操作
- 這些操作如何并行執行
- 如何創建控制依賴性來優化并行執行
是時候將計算分布在多個服務器上了!
## 多個服務器的多個設備
要跨多臺服務器運行圖形,首先需要定義一個集群。 一個集群由一個或多個 TensorFlow 服務器組成,稱為任務,通常分布在多臺機器上(見圖 12-6)。 每項任務都屬于一項作業。 作業只是一組通常具有共同作用的任務,例如跟蹤模型參數(例如,參數服務器通常命名為`"ps"`,parameter server)或執行計算(這樣的作業通常被命名為`"worker"`)。

以下集群規范定義了兩個作業`"ps"`和`"worker"`,分別包含一個任務和兩個任務。 在這個例子中,機器A托管著兩個 TensorFlow 服務器(即任務),監聽不同的端口:一個是`"ps"`作業的一部分,另一個是`"worker"`作業的一部分。 機器B僅托管一臺 TensorFlow 服務器,這是`"worker"`作業的一部分。
```py
cluster_spec = tf.train.ClusterSpec({
"ps": [
"machine-a.example.com:2221", # /job:ps/task:0
],
"worker": [
"machine-a.example.com:2222", # /job:worker/task:0
"machine-b.example.com:2222", # /job:worker/task:1
]})
```
要啟動 TensorFlow 服務器,您必須創建一個服務器對象,并向其傳遞集群規范(以便它可以與其他服務器通信)以及它自己的作業名稱和任務編號。 例如,要啟動第一個輔助任務,您需要在機器 A 上運行以下代碼:
```py
server = tf.train.Server(cluster_spec, job_name="worker", task_index=0)
```
每臺機器只運行一個任務通常比較簡單,但前面的例子表明 TensorFlow 允許您在同一臺機器上運行多個任務(如果需要的話)。 如果您在一臺機器上安裝了多臺服務器,則需要確保它們不會全部嘗試抓取每個 GPU 的所有 RAM,如前所述。 例如,在圖12-6中,`"ps"`任務沒有看到 GPU 設備,想必其進程是使用`CUDA_VISIBLE_DEVICES =""`啟動的。 請注意,CPU由位于同一臺計算機上的所有任務共享。
如果您希望進程除了運行 TensorFlow 服務器之外什么都不做,您可以通過告訴它等待服務器使用`join()`方法來完成,從而阻塞主線程(否則服務器將在您的主線程退出)。 由于目前沒有辦法阻止服務器,這實際上會永遠阻止:
```py
server.join() # blocks until the server stops (i.e., never)
```
## 開始一個會話
一旦所有任務啟動并運行(但還什么都沒做),您可以從位于任何機器上的任何進程(甚至是運行中的進程)中的客戶機上的任何服務器上打開會話,并使用該會話像普通的本地會議一樣。比如:
```py
a = tf.constant(1.0)
b = a + 2
c = a * 3
with tf.Session("grpc://machine-b.example.com:2222") as sess:
print(c.eval()) # 9.0
```
這個客戶端代碼首先創建一個簡單的圖形,然后在位于機器 B(我們稱之為主機)上的 TensorFlow 服務器上打開一個會話,并指示它求值`c`。 主設備首先將操作放在適當的設備上。 在這個例子中,因為我們沒有在任何設備上進行任何操作,所以主設備只將它們全部放在它自己的默認設備上 - 在這種情況下是機器 B 的 GPU 設備。 然后它只是按照客戶的指示求值`c`,并返回結果。
## 主機和輔助服務
客戶端使用 gRPC 協議(Google Remote Procedure Call)與服務器進行通信。 這是一個高效的開源框架,可以調用遠程函數,并通過各種平臺和語言獲取它們的輸出。它基于 HTTP2,打開一個連接并在整個會話期間保持打開狀態,一旦建立連接就可以進行高效的雙向通信。
數據以協議緩沖區的形式傳輸,這是另一種開源 Google 技術。 這是一種輕量級的二進制數據交換格式。
TensorFlow 集群中的所有服務器都可能與集群中的任何其他服務器通信,因此請確保在防火墻上打開適當的端口。
每臺 TensorFlow 服務器都提供兩種服務:主服務和輔助服務。 主服務允許客戶打開會話并使用它們來運行圖形。 它協調跨任務的計算,依靠輔助服務實際執行其他任務的計算并獲得結果。
## 固定任務的操作
通過指定作業名稱,任務索引,設備類型和設備索引,可以使用設備塊來鎖定由任何任務管理的任何設備上的操作。 例如,以下代碼將`a`固定在`"ps"`作業(即機器 A 上的 CPU)中第一個任務的 CPU,并將`b`固定在`"worker"`作業的第一個任務管理的第二個 GPU (這是 A 機上的 GPU#1)。 最后,`c`沒有固定在任何設備上,所以主設備將它放在它自己的默認設備上(機器 B 的 GPU#0 設備)。
```py
with tf.device("/job:ps/task:0/cpu:0")
a = tf.constant(1.0)
with tf.device("/job:worker/task:0/gpu:1")
b = a + 2
c = a + b
```
如前所述,如果您省略設備類型和索引,則 TensorFlow 將默認為該任務的默認設備; 例如,將操作固定到`"/job:ps/task:0"`會將其放置在`"ps"`作業(機器 A 的 CPU)的第一個任務的默認設備上。 如果您還省略了任務索引(例如,`"/job:ps"`),則 TensorFlow 默認為`"/task:0"`。如果省略作業名稱和任務索引,則 TensorFlow 默認為會話的主任務。
## 跨多個參數服務器的分片變量
正如我們很快會看到的那樣,在分布式設置上訓練神經網絡時,常見模式是將模型參數存儲在一組參數服務器上(即`"ps"`作業中的任務),而其他任務則集中在計算上(即 ,`"worker"`工作中的任務)。 對于具有數百萬參數的大型模型,在多個參數服務器上分割這些參數非常有用,可以降低飽和單個參數服務器網卡的風險。 如果您要將每個變量手動固定到不同的參數服務器,那將非常繁瑣。 幸運的是,TensorFlow 提供了`replica_device_setter()`函數,它以循環方式在所有`"ps"`任務中分配變量。 例如,以下代碼將五個變量引入兩個參數服務器:
```py
with tf.device(tf.train.replica_device_setter(ps_tasks=2):
v1 = tf.Variable(1.0) # pinned to /job:ps/task:0
v2 = tf.Variable(2.0) # pinned to /job:ps/task:1
v3 = tf.Variable(3.0) # pinned to /job:ps/task:0
v4 = tf.Variable(4.0) # pinned to /job:ps/task:1
v5 = tf.Variable(5.0) # pinned to /job:ps/task:0
```
您不必傳遞`ps_tasks`的數量,您可以傳遞集群`spec = cluster_spec`,TensorFlow 將簡單計算`"ps"`作業中的任務數。
如果您在塊中創建其他操作,則不僅僅是變量,TensorFlow 會自動將它們連接到`"/job:worker"`,默認為第一個由`"worker"`作業中第一個任務管理的設備。 您可以通過設置`worker_device`參數將它們固定到其他設備,但更好的方法是使用嵌入式設備塊。 內部設備塊可以覆蓋在外部塊中定義的作業,任務或設備。 例如:
```py
with tf.device(tf.train.replica_device_setter(ps_tasks=2)):
v1 = tf.Variable(1.0) # pinned to /job:ps/task:0 (+ defaults to /cpu:0)
v2 = tf.Variable(2.0) # pinned to /job:ps/task:1 (+ defaults to /cpu:0)
v3 = tf.Variable(3.0) # pinned to /job:ps/task:0 (+ defaults to /cpu:0)
[...]
s = v1 + v2 # pinned to /job:worker (+ defaults to task:0/gpu:0)
with tf.device("/gpu:1"):
p1 = 2 * s # pinned to /job:worker/gpu:1 (+ defaults to /task:0)
with tf.device("/task:1"):
p2 = 3 * s # pinned to /job:worker/task:1/gpu:1
```
這個例子假設參數服務器是純 CPU 的,這通常是這種情況,因為它們只需要存儲和傳送參數,而不是執行密集計算。
(未完成)