《Scikit-Learn與TensorFlow機器學習實用指南》 第12章 設備和服務器上的分佈式 TensorFlow



在第 11 章,我們討論了幾種可以明顯加速訓練的技術:更好的權重初始化,批量標準化,複雜的優化器等等。 但是,即使採用了所有這些技術,在具有單個 CPU 的單臺機器上訓練大型神經網絡可能需要幾天甚至幾周的時間。

在本章中,我們將看到如何使用 TensorFlow 在多個設備(CPU 和 GPU)上分配計算並將它們並行運行(參見圖 12-1)。 首先,我們會先在一臺機器上的多個設備上分配計算,然後在多臺機器上的多個設備上分配計算。

圖12-1 在多臺設備上並行執行TensorFlow圖

與其他神經網絡框架相比,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 的博客文章中,提供了在 Amazon AWS GPU 實例上使用 Python 3.5 設置 TensorFlow 0.9 的詳細說明。將它更新到最新版本的 TensorFlow 應該不會太難。 Google 還發布了一項名爲 Cloud Machine Learning 的雲服務來運行 TensorFlow 圖表。 2016 年 5 月,他們宣佈他們的平臺現在包括配備張量處理器(TPU)的服務器,專門用於機器學習的處理器,比許多 GPU 處理 ML 任務要快得多。 當然,另一種選擇只是購買你自己的 GPU 卡。 Tim Dettmers 寫了一篇很棒的博客文章(http://timdettmers.com/2018/11/05/which-gpu-for-deep-learning/)來幫助你選擇,他會定期更新它。

您必須下載並安裝相應版本的 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)。

圖12-2 TensorFlow使用CUDA和cuDNN控制GPU,從而加速訓練DNN

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

$ nvidia-smi
Wed Sep 16 09:50:03 2016
+------------------------------------------------------+
| NVIDIA-SMI 352.63     Driver Version: 352.63         |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GRID K520           Off  | 0000:00:03.0     Off |                  N/A |
| N/A   27C    P8    17W / 125W |     11MiB /  4095MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

最後,您必須安裝支持 GPU 的 TensorFlow。 如果你使用virtualenv創建了一個獨立的環境,你首先需要激活它:

$ cd $ML_PATH               # Your ML working directory (e.g., $HOME/ml)
$ source env/bin/activate

然後安裝合適的支持 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)。

圖12-3 每個程序都有兩塊GPU

另一種選擇是告訴 TensorFlow 只抓取一小部分內存。 例如,要使 TensorFlow 只佔用每個 GPU 內存的 40%,您必須創建一個ConfigProto對象,將其gpu_options.per_process_gpu_memory_fraction選項設置爲 0.4,並使用以下配置創建session

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。

圖12-4 每個程序都可以使用四塊GPU,但每個程序只分配了40%的RAM

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

$ nvidia-smi
[...]
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|    0      5231    C   python                                        1677MiB |
|    0      5262    C   python                                        1677MiB |
|    1      5231    C   python                                        1677MiB |
|    1      5262    C   python                                        1677MiB |
[...]

另一種選擇是告訴 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不固定在任何設備上,因此將放置在默認設備上:

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"

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

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 並行執行TensorFlow計算圖

例如,在圖 12-5 中,操作ABC是源操作,因此可以立即進行求值。 操作AB放置在 GPU#0 上,因此它們被髮送到該設備的內部線程池,並立即進行並行求值。 操作A正好有一個多線程內核; 它的計算被分成三部分,這些部分由內部線程池並行執行。 操作C轉到 GPU#1 的內部線程池。

一旦操作C完成,操作DE的依賴性計數器將遞減並且都將達到 0,因此這兩個操作將被髮送到操作內線程池以執行。

您可以通過設置inter_op_parallelism_threads選項來控制內部線程池的線程數。 請注意,您開始的第一個會話將創建內部線程池。 除非您將use_per_session_threads選項設置爲True,否則所有其他會話都將重用它們。 您可以通過設置intra_op_parallelism_threads選項來控制每個內部線程池的線程數。

控制依賴關係

在某些情況下,即使所有依賴的操作都已執行,推遲對操作的求值可能也是明智之舉。例如,如果它使用大量內存,但在圖形中只需要更多內存,則最好在最後一刻對其進行求值,以避免不必要地佔用其他操作可能需要的 RAM。 另一個例子是依賴位於設備外部的數據的一組操作。 如果它們全部同時運行,它們可能會使設備的通信帶寬達到飽和,並最終導致所有等待 I/O。 其他需要傳遞數據的操作也將被阻止。 順序執行這些通信繁重的操作將是比較好的,這樣允許設備並行執行其他操作。

推遲對某些節點的求值,一個簡單的解決方案是添加控制依賴關係。 例如,下面的代碼告訴 TensorFlow 僅在求值完ab之後才求值xy

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依賴於xy,所以求值z也意味着等待ab進行求值,即使它並未顯式存在於control_dependencies()塊中。 此外,由於b依賴於a,所以我們可以通過在[b]而不是[a,b]上創建控制依賴關係來簡化前面的代碼,但在某些情況下,“顯式比隱式更好”。

很好!現在你知道了:

  • 如何以任何您喜歡的方式在多個設備上進行操作
  • 這些操作如何並行執行
  • 如何創建控制依賴性來優化並行執行

是時候將計算分佈在多個服務器上了!

多個服務器的多個設備

要跨多臺服務器運行圖形,首先需要定義一個集羣。 一個集羣由一個或多個 TensorFlow 服務器組成,稱爲任務,通常分佈在多臺機器上(見圖 12-6)。 每項任務都屬於一項作業。 作業只是一組通常具有共同作用的任務,例如跟蹤模型參數(例如,參數服務器通常命名爲"ps",parameter server)或執行計算(這樣的作業通常被命名爲"worker")。

圖12-6 TensorFlow集羣

以下集羣規範定義了兩個作業"ps""worker",分別包含一個任務和兩個任務。 在這個例子中,機器A託管着兩個 TensorFlow 服務器(即任務),監聽不同的端口:一個是"ps"作業的一部分,另一個是"worker"作業的一部分。 機器B僅託管一臺 TensorFlow 服務器,這是"worker"作業的一部分。

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 上運行以下代碼:

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()方法來完成,從而阻塞主線程(否則服務器將在您的主線程退出)。 由於目前沒有辦法阻止服務器,這實際上會永遠阻止:

server.join()  # blocks until the server stops (i.e., never) 

開始一個會話

一旦所有任務啓動並運行(但還什麼都沒做),您可以從位於任何機器上的任何進程(甚至是運行中的進程)中的客戶機上的任何服務器上打開會話,並使用該會話像普通的本地會議一樣。比如:

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 設備)。

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"任務中分配變量。 例如,以下代碼將五個變量引入兩個參數服務器:

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參數將它們固定到其他設備,但更好的方法是使用嵌入式設備塊。 內部設備塊可以覆蓋在外部塊中定義的作業,任務或設備。 例如:

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 的,這通常是這種情況,因爲它們只需要存儲和傳送參數,而不是執行密集計算。

(負責這章的人沒有全部完成,等整理完最後一章,我再來補上)



發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章