[翻譯] TensorFlow 分散式之論文篇 “TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems”

[翻譯] TensorFlow 分散式之論文篇 “TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems”

本系列我們開始分析 TensorFlow 的分散式。之前在機器學習分散式這一系列分析之中,我們大多是以 PyTorch 為例,結合其他框架/庫來穿插完成。但是缺少了 TensorFlow 就會覺得整個世界(系列)都是不完美的,不單單因為 TensorFlow 本身的影響力,更因為 TensorFlow 分散式有自己的鮮明特色,對於技術愛好者來說是一個巨大寶藏。

讀論文有一種原則是:本領域最經典的論文,近5年最熱的論文,近1年最新的論文。按照這個原則,本文主要介紹一篇 TensorFlow 經典論文 TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems。大家如果讀了下面論文就會發現 TensorFlow分散式的博大精深。

本文圖來自原始論文。

1. 原文摘要

基於我們在 DistBelief 方面的經驗,以及對期望的系統特性和訓練/使用神經網路需求的更全面的理解,我們構建了 TensorFlow ,這是我們的第二代系統,用於實現和部署大規模機器學習模型。 TensorFlow 採用類似數據流的模型來描述計算,並將其映射到各種不同的硬體平台上,從在 Android 和 iOS 等移動設備平台上運行推理,到使用包含一個或多個 GPU 卡的單機的中等規模訓練和推理系統,再到在數百台具有數千個 GPU 的專用主機上運行的大規模訓練系統。

擁有一個能夠跨越如此廣泛的平台的單一系統可以顯著簡化機器學習系統的實際使用,因為我們發現,如果大規模訓練和小規模部署都分別有自己的獨立系統,則會導致巨大的維護負擔和漏洞。 TensorFlow 計算被表示為有狀態數據流圖,我們致力於使系統具有足夠的靈活性,以便用戶可以快速試驗新模型,系統同時也具有足夠高的性能和魯棒性,可以被用於機器學習模型的訓練和部署。

為了將神經網路訓練擴展到更大規模的部署, TensorFlow 允許客戶機通過複製和並行執行核心模型數據流圖來輕鬆表達各種並行性,這樣可以使用許多不同的計算設備來更新一組共享參數或其他共享狀態。對計算描述的適度更改允許用戶實現多種不同的並行方法。

TensorFlow 允許在參數更新的一致性方面具有一定的靈活性,這些寬鬆的同步要求允許我們可以在一些較大的部署中更加輕鬆。與 DistBelief 相比, TensorFlow 的編程模型更加靈活,性能顯著提高,支援在更廣泛的異構硬體平台上訓練和使用更廣泛的模型。

2. 編程模型和基本概念

TensorFlow 計算由一組節點組成的有向來描述。該圖表示一個數據流計算,也允許讓某些類型的節點來維護和更新持久狀態,並以類似於 Naiad 的方式在圖中實現分支和循環控制結構。客戶端通常使用一種前端語言(C++或Python)構建計算圖。下圖顯示了使用 Python 前端構建並執行 TensorFlow 計算圖的示例片段,

圖 1. TensorFlow 計算圖示例片段

圖 2,計算圖

在 TensorFlow 圖中,每個節點表示操作的實例,其具有零個或多個輸入和零個或多個輸出。在計算圖中沿普通邊流動的值(從輸出到輸入)被稱為張量。張量是任意維數組,其基本元素類型在計算圖構造時被指定或推斷出來。特殊的邊(被稱為控制依賴關係)也可以存在於圖中:這些邊上沒有數據流,但它們表明控制依賴關係的源節點必須在控制依賴關係的目標節點開始執行之前完成執行。因為我們的模型包含可變狀態,因此客戶可以直接使用控制依賴關係來強制執行。我們的實現有時也會在內部插入控制依賴項,以強制某些獨立操作之間的順序,例如,可以控制峰值記憶體使用。

2.1 運算元(Operations)與核(Kernels)

運算元(Operation)表示一個抽象計算(例如,”矩陣乘法”或”加法”)。一個運算元可以擁有屬性,但是所有屬性必須在計算圖構造時被提供或推斷出來,這樣才能實例化一個執行該運算元的節點。屬性的一個常見用途是使運算元在不同的張量元素類型上具有多態性(例如,加法運算元即支援兩個類型為 float 的 tensors 相加,也支援兩個類型為 int32的張量相加)。

核(Kernel)是可以在特定類型的設備(例如CPU或GPU)上運行的運算元的具體實現。 TensorFlow 通過註冊機制定義了一系列運算元和核,這樣意味著用戶可以通過鏈接其他運算元和/或內核來進行擴展。下圖顯示了 TensorFlow 庫中內置的一些運算元。

表 1. 運算元

2.2 會話(Sessions)

客戶端程式通過創建會話與 TensorFlow 系統交互。會話先構建一個空計算圖,為了創建計算圖,會話介面支援Extend 方法,該方法可以讓客戶端附加節點和邊來擴充這個空圖,然後進行計算。

2.3 變數(Variables)

在大多數計算中,一個圖會被執行多次,而大多數張量在圖的單次執行之後都不會存在。然而,有一些張量是在計算圖的執行過程之中始終存在的,位置也是固定的,其不能正常流動但是可以更新,比如模型的參數,這就引出了變數這個概念。

變數是一種特殊的操作,它返回持久可變張量的句柄,這些句柄可以被傳遞給少量特殊的操作,例如 AssignAssignAdd(相當於+=),通過這些操作就可以改變這些變數引用的張量。

3. 實現

TensorFlow 系統中的主要組件是客戶端,它使用會話介面與主機以及一個或多個工作進程進行通訊。每個工作進程負責協調對一個或多個計算設備(如 CPU 內核或 GPU 卡)的訪問以及按照主設備的指示在這些設備上執行計算圖節點。 TensorFlow 介面有本地分散式實現兩種。

  • 當客戶端、master 和 worker 都在單個機器上單個進程的上下文之中運行時(如果機器安裝了多個 GPU 卡,則可能使用多個設備),將使用本地實現。

  • 分散式實現與本地實現共享大部分程式碼,但支援通過一個環境對其進行擴展,在該環境中,客戶端、master和 worker 都可以在不同機器上不同的進程中。在我們的分散式環境中,這些不同的任務是 job 之中的一些容器,這些 job 由集群調度系統來管理。這兩種不同的模式如下圖所示。

圖 3. 調度模式

3.1 設備(Devices)

設備是 TensorFlow 的計算核心。每個工作者負責一個或多個設備,每個設備都有一個設備類型和名稱。設備名稱由以下幾部分組成:

  • 設備類型。
  • 設備在工作者中的索引。
  • 分散式設置中對於工作者所在作業和任務的標識(如果設備是進程本地的,則為 localhost)。

示例設備名稱的樣例如下:

  • “/job:localhost/device:cpu:0″。
  • “/job:worker/task:17/device:gpu:3″。

PyTorch 有針對 CPU 和 GPU 的設備介面的實現,其他設備類型可以通過註冊機制提供新設備實現。每個設備對象負責管理設備記憶體的分配和釋放,以及執行 TensorFlow 下發的核方法。

3.2 張量

在我們的實現中,張量是一個類型化的多維數組。我們支援多種張量元素類型,包括大小從 8 位到 64 位的有符號和無符號整數、IEEE 浮點和雙精度類型、複數類型和字元串類型(arbitrary byte array)。張量所在設備的分配器負責管理張量的存儲區,張量存儲緩衝區是引用計數的,在沒有引用保留時會進行釋放。

3.3 單設備執行

讓我們首先考慮最簡單的執行場景:一個擁有單個設備的工作者進程。計算圖中的節點按照節點之間依賴關係的順序來執行。我們將跟蹤每個節點尚未執行的依賴項數量的計數。一旦此計數降至零,該節點就有資格執行,並被添加到就緒隊列中。就緒隊列以某種未指定的順序進行處理,其將節點的核方法執行委託給設備對象。當節點完成執行時,依賴於此已完成節點的所有節點的計數都將減少。

3.4 多設備執行

一旦一個系統有多個設備,就有兩個主要的複雜問題:如何決定將每個節點的計算放在哪個設備上,如何管理這些放置(Placement )所帶來的跨設備數據通訊。本小節討論這兩個問題。

3.4.1 決定設備(Node Placement)

給定計算圖之後, TensorFlow 實現的主要職責之一是將計算映射到可用設備集。本文給出了該演算法的一個簡化版本。

布局(placement)演算法的一個輸入是成本模型(cost model),該模型包含每個計算圖節點的輸入和輸出張量的大小(位元組)估計,以及每個節點在獲得輸入張量之後所需的計算時間估計。該成本模型要麼基於與不同操作類型相關的啟發式靜態估計,要麼基於計算圖早期執行的實際布局決策集進行測量/決定。

布局演算法首先運行計算圖的模擬執行,然後使用貪婪啟發式為圖中的每個節點選擇一個設備。此模擬生成的”節點到設備放置關係”最終也用作實際執行的放置。布局演算法從計算圖的源開始,並在前進過程中模擬系統中每個設備上的活動,在此遍歷中:

  • 當到達了一個節點,就考慮此節點的可使用設備集(如果設備不提供用戶希望的實現特定操作的內核,則設備就不使用)。
  • 對於具有多個可用設備的節點,布局演算法使用貪婪啟發式演算法,看看將節點放置在每個可能設備上對節點完成時間會造成怎樣的影響。該啟發式方法不僅考慮了在成本模型中,該類設備上操作的估計或測量執行時間,還包括為了將輸入從其他設備傳輸到該節點而引入的通訊成本。然後選擇最快完成節點操作的設備作為該操作的設備,
  • 布局(placement)過程繼續進行,以便為圖中的其他節點(包括現在準備好模擬執行的下游節點)做出放置決策。

3.4.2 跨設備通訊(Cross-Device Communication)

一旦決定了節點如何放置到設備之上(node placement),圖就被劃分成一組子圖,每個設備一個子圖。任何跨設備的從 x 到 y 的邊將被刪除,並替換為兩條新邊:

  • 一條邊是在 x 子圖中,從 x 到新的 Send 節點。
  • 一條邊是在 y 的子圖之中,從對應的 Receive 節點到 y。

請參見下圖。

圖 4 插入發送/接收節點之前和之後

在運行時,發送和接收節點將會彼此協調如何在設備之間傳輸數據。這使我們能夠把發送和接收的所有通訊隔離出來,從而簡化運行時(runtime)的其餘部分。

當我們插入發送和接收節點時,我們規範如下:特定設備上特定張量的所有用戶都使用同一個接收節點,而不是特定設備上的每個下游用戶都擁有一個自己的接收節點。這確保了所需張量的數據在”源設備→ 目標設備對”之間只傳輸一次,並且目標設備上張量的記憶體只分配一次,而不是多次(例如,參見上圖中的節點 b 和 c)。

通過以這種方式處理通訊,我們還允許將不同設備上的圖的各個節點的調度分散到工作者之中:發送和接收節點在不同的工作者和設備之間傳遞必要的同步,這樣就把主節點從調度任務之中解放出來。主節點只需要向每個具有計算圖的任何節點的工作者發出單個 Run 請求(每次計算圖執行),而不需要參與每個節點或每個跨設備通訊的調度。這使得系統更具可伸縮性可擴展性,並且和主節點強制執行調度相比,可以執行更細粒度的節點執行策略。

3.5 分散式執行

計算圖的分散式執行與多設備執行非常相似。在決定設備如何放置之後,將為每個設備創建一個子圖。發送/接收節點對在跨工作進程通訊時候使用遠程通訊機制(如 TCP 或 RDMA)來跨機器邊界移動數據。

3.5.1 容錯(Fault Tolerance)

我們可以在許多地方檢測到分散式執行中的故障,目前主要依靠如下兩種方式:

  • 檢測到發送和接收節點對之間的通訊錯誤。
  • 從主進程到每個工作進程的定期健康檢查。

當檢測到故障時,整個計算圖執行將中止並從頭開始。因為變數(Variable)節點指的是在圖的執行過程中持續存在的張量,所以我們支援設置一致性檢查點,以此來在重新啟動時恢復這些狀態。具體來說,每個變數節點都連接到一個 Save 節點。這些保存節點定期執行,例如每 N 次迭代執行一次,或每 N 秒執行一次。當它們執行時,變數的內容被寫入持久存儲,例如分散式文件系統。類似地,每個變數也都連接到一個恢復節點,該節點僅在重新啟動後的第一次迭代中啟用。

4. 擴展(Extensions)

在本節中,我們將介紹基本編程模型的幾個更高級的特性。

4.1 計算梯度

許多優化演算法,包括常見的機器學習訓練演算法(如隨機梯度下降),會使用一組輸入來計算一個成本函數(cost function)的梯度。因為這是一種常見的需求,所以 TensorFlow 內置了對自動梯度計算的支援。如果一個 TensorFlow 計算圖中的張量 C 可能通過一個複雜的操作子圖依賴於一組張量{\(X_k\)},那麼一個內置函數將返回張量集{\(dC/dX_k\)}。梯度張量和其他張量一樣,通過使用以下步驟擴展 TensorFlow 圖來計算。

張量 C 依賴於張量 I,當 TensorFlow 需要計算張量 C 相對於張量I的梯度時,它首先在計算圖中找到從 I 到 C 的路徑。然後它從 C 回溯到 I,對於反向路徑上的每個操作,它會向 TensorFlow 圖添加一個節點,使用鏈式規則沿向後路徑合成偏導數。

新添加的節點為前向路徑中的相應操作計算”梯度函數”。梯度函數可以通過任何操作註冊。該函數不僅將沿反向路徑計算的部分梯度作為輸入,還可以選擇正向操作的輸入和輸出。圖5顯示了根據圖2示例計算的成本梯度。灰色箭頭顯示梯度函數的潛在輸入,該函數不用於所示的特定操作。

圖 5 計算這些梯度所需的附加值為:

\[[db,dW,dx] = tf.gradients(C,[b,W,x])
\]

自動梯度計算會使優化更加複雜化,尤其是記憶體使用。當執行前向計運算元圖(即用戶顯式構造的子圖)時,啟發式優化演算法可以通過觀察計算圖的構造順序來決定下一步執行哪個節點,這通常意味著臨時輸出在構建後很快就會被佔用,因此它們的記憶體可以很快被重用。

當啟發式無效時,用戶可以通過更改計算圖構造的順序,或添加控制依賴項來優化記憶體使用。但是,當梯度節點自動添加到計算圖中時,用戶的控制能力會降低,啟發式演算法可能會崩潰。特別是,因為梯度反轉了正向計算順序,因此在計算圖執行中,早期使用的張量在梯度計算的末尾經常再次需要。這種張量會佔用大量稀缺的 GPU 記憶體,從而不必要地限制計算量。我們正在積極改進記憶體管理,以便更好地處理此類情況。選項包括使用更複雜的啟發演算法來確定計算圖執行的順序,重新計算張量而不是將其保留在記憶體中,以及將長期張量從 GPU 記憶體交換到更大的主機 CPU 記憶體。

4.2 局部執行(Partial Execution)

客戶機通常只想執行整個執行圖的一個子圖。為了支援這一點,一旦客戶機在會話中設置了計算圖,我們的 Run 方法允許客戶機執行整個圖的任意子圖,並沿圖中的任意邊輸入任意數據,以及沿圖中任意邊獲取數據。

圖中的每個節點都有一個名稱,節點的每個輸出由源節點名稱和節點的輸出埠標識,從 0 開始編號(例如,”bar:0″ 表示”bar”節點的第一個輸出,而”bar:1″表示第二個輸出)。

Run 調用的兩個參數有助於定義將要執行的計算圖的確切子圖。首先,Run 調用接受類型為 name:port 名稱到”fed”張量值的映射作為輸入。其次,Run 調用接受 output names,這是一個輸出 name[:port] 列表,其指定了應執行哪些節點。如果名稱中存在埠部分,則如果 Run 調用成功完成,應將節點的特定輸出張量值返回給客戶端。

計算圖可以基於輸入和輸出的值進行轉換。輸入中每個 node:port 都替換為一個 feed 節點,該節點將從用於 Run 調用的 Rendezvous 對象中獲取輸入張量。類似地,每個帶有埠的輸出名稱都連接到一個特殊的 fetch 節點,該節點被用來保存輸出張量,並在運行調用完成時將其返回給客戶端。最後,一旦通過插入這些特殊的 feed 和 fetch 節點重寫了計算圖,要執行的節點集可以通過以下方式確定:從每個由輸出指定的節點開始,使用圖依賴關係在圖中進行後向傳播,以確定為了計算輸出而必須在重寫圖中執行的完整節點集。下圖在左側顯示了一個原始圖,以及在使用 inputs{b} 和 outputs{f:0} 調用 Run 時生成的轉換圖。因為我們只需要計算節點f的輸出,所以我們不會執行節點 d 和 e,因為它們與 f 的輸出沒有任何關係。

圖 6 局部執行前後

4.3 設備約束(Device Constraints)

TensorFlow 客戶端可以通過為節點提供部分約束來控制節點在設備上的位置,這些約束與節點可以在哪些設備上執行有關。例如,“僅將此節點放置在 GPU 類型的設備上”,或“此節點可以放置在/job:worker/task:17“中的任何設備上,或“將此節點與名為variable13“的節點合併”。在這些約束條件的約束範圍內,布局演算法(placement algorithm)負責完成節點到設備的分配,以提供快速的計算執行,並滿足設備自身施加的各種約束,例如,限制設備上執行其計算圖節點子集所需的記憶體總量。

支援此類約束要求更改前面描述的布局演算法。我們首先計算每個節點的可行設備集,然後在共定位約束圖(graph of colocation constraints)上使用 union find 來計算出必須放置在一起的圖組件。對於每個這樣的組件,我們計算可行設備集的交集。

4.4 控制流

雖然沒有任何顯式控制流的數據流圖也非常有表達能力,但我們發現,在很多情況下,如果支援條件和循環,則可以用更簡潔和有效來表示機器學習演算法。

與 Arvind 描述的數據流機(dataflow-machine)方法一樣,我們在 TensorFlow 中引入了一小部分控制流原語操作符,並將 TensorFlow 推廣到可以處理循環數據流圖。SwitchMerge 運算符允許我們根據布爾張量的值來跳過整個子圖的執行。Enter,LeaveNextIteration 運算符允許我們表示迭代。高層級的編程結構,如 if-conditionals 和 while-loops 則可以使用這些控制流操作符來輕鬆地編譯成數據流計算圖。

TensorFlow 運行時實現了標籤(tags)和幀(frames)的概念,其在概念上類似於 MIT 標記令牌機(MIT Tagged-Token machine)。循環的每個迭代都由一個 tag 唯一標識,其執行狀態由一個 frame 表示。只要輸入準備好,它就可以進入迭代,因此可以同時執行多個迭代。

如何為分散式系統處理循環控制的狀態? TensorFlow 使用分散式協調機制來執行帶有控制流的圖。通常,循環可以包含分配給許多不同設備的節點。因此,管理循環的狀態成為分散式終止檢測問題。TensorFlow 的解決方案是基於圖重寫。在圖分區過程中,我們自動向每個分區添加控制節點。這些節點實現了一個小型狀態機,它協調每個迭代的開始和結束,並決定最終循環的結束。對於每個迭代,擁有循環終止斷言(predicate)的設備向每個參與的設備發送一條控制消息。

如上所述,我們通常通過梯度下降來訓練機器學習模型,並將梯度計算表示為數據流圖的一部分。當模型包含控制流操作時,我們必須在相應的梯度計算中考慮它們。例如,對於具有 if-conditional 的模型,梯度計算需要知道採用了條件的哪個分支,然後將梯度邏輯應用於該分支。類似地,帶有 while-loop 的模型的梯度計算需要知道進行了多少次迭代,並且還將依賴於在這些迭代期間計算的中間值。目前依然採用重寫計算圖技術來記錄梯度計算所需的值。

4.5 輸入操作

雖然可以通過 feed 節點把輸入數據提供給計算調用,但用於訓練大規模機器學習模型的另一種常見機制是在圖中部署有特定的輸入操作節點,這種節點通常配置成一組文件名,該節點每次執行時產生一個張量,該張量包含存儲在該組文件中的數據的一個或多個樣本。這樣就允許將數據直接從底層存儲系統讀取到電腦記憶體中,然後電腦記憶體將對數據執行後續處理。在客戶端進程與工作進程分開的配置中,如果數據被饋送,則通常需要額外的網路躍點 hop(從存儲系統到客戶端,然後從客戶端到工作進程,而不是使用輸入節點時直接從存儲系統傳輸到工作進程)。

4.6 隊列

隊列是我們添加到 TensorFlow 中的一個有用特性。它們允許計算圖的不同部分進行非同步操作,並通過入隊(Enqueue)和出隊(Dequeue)操作傳遞數據。入隊操作可以阻塞,直到隊列中有可用的空間,而出隊操作也可以阻塞,直到隊列中有所需的最少數量的元素可用。隊列的一種用途是,當機器學習模型的計算部分仍在處理前一批數據時,模型可以從磁碟文件中預取輸入數據。它們也可用於其他類型的分組操作,包括累積多個梯度,這樣可以把小 batch 組合成為一個大 batch,以便在大的批次上計算更複雜的梯度組合,或將循環語言模型的不同輸入句子分組到大致相同長度的句子箱(bin)中,然後可以更有效地處理這些問題。

除了普通的 FIFO 隊列之外,我們還實現了一個 shuffling 隊列,它在一個大的記憶體緩衝區內隨機 shuffle 其元素。例如,對於希望用隨機順序來處理樣本的機器學習演算法來說,這種洗牌功能非常有用。

4.7 容器

容器是 TensorFlow 中用於管理長期可變狀態的機制。默認容器將會一直持續到進程終止,但我們也允許使用其他的命名容器。容器存儲變數的備份,可以通過完全清除容器中的內容來重置容器。通過使用容器可以在不同會話的完全不相交的計算圖之間共享狀態。

5 優化

在本節中,我們將介紹 TensorFlow 實現中的一些優化,這些優化可以提高系統的性能或資源利用率。

5.1 消除公共子表達式

由於計算圖的構造通常由客戶機程式碼中的許多不同層的抽象來完成,因此計算圖很容易存在重複計算。為了處理這個問題,我們實現了類似於Click(原始論文參考文獻12)描述的演算法,該演算法在計算圖上運行,並將具有相同輸入和操作類型的多個操作副本縮減到其中的一個節點,並會把相應的邊進行重定向。

5.2 控制數據傳輸和記憶體使用

仔細安排 TensorFlow 操作可以提高系統的性能,特別是在數據傳輸和記憶體使用方面。具體而言,調度可以減少中間結果保存在記憶體中的時間,從而減少記憶體消耗峰值。這種減少對於記憶體不足的 GPU 設備尤為重要。此外,重新安排設備間的數據通訊可以減少對網路資源的爭奪。

雖然有很多機會可以安排優化,但這裡我們重點關注一個我們認為特別必要和有效的機會。它涉及接收節點讀取遠程值的計劃。如果不採取預防措施,這些節點可能會比必要時啟動得更快,可能在執行開始時一次啟動。通過執行運籌學中常見的儘快/儘可能晚(ASAP/ALAP)計算,我們分析圖的關鍵路徑以估計何時啟動接收節點。然後,我們插入控制邊,目的是在需要這些節點的結果之前取消這些節點的起始位置。

5.3 非同步核

除了正常同步內核之外,我們的框架還支援非阻塞內核。這樣的非阻塞內核使用一個稍有不同的介面,通過該介面,一個 continuation 被傳遞給計算方法,該 continuation 應該在內核執行完成時調用。這是一種針對具有多個活動執行緒環境的優化。在這些環境中,記憶體使用或其他資源相對昂貴,因此我們應該避免在等待 I/O 或其他事件發生時長期佔用執行執行緒。非同步內核的示例包括接收內核、排隊內核和出列內核(如果隊列空間不可用或沒有可讀取的數據,則可能需要分別進行阻塞)。

5.4 優化函數庫

目前有很多高度優化的函數庫,比如 BLAS 和 cuBLAS,另外,cuda convnet 和 cuDNN 可以在 GPU 之上處理卷積核。因此我們利用這些函數庫實現了系統的很多內核。

5.5 有損壓縮(Lossy Compression)

一些機器學習演算法(比如通常用於訓練神經網路的演算法)具有抗雜訊和降低演算法精度的能力。在設備之間發送數據時(有時在同一台機器內的設備之間),我們通常使用高精度的有損壓縮,這一方式類似於分散式系統。例如,我們通常插入特殊的轉換節點,將 32 位浮點表示轉換為 16 位浮點表示(不是 IEEE 16 位浮點標準,而是 32 位 IEEE 794 浮點格式,但尾數中的精度降低了 16 位),然後在通訊信道的另一端轉換回 32 位表示(只需為尾數的丟失部分填入零,因為這樣在 32→ 16→ 32 位轉換時計算花費更少)。

6. 編程技巧

上面講的都是一些系統級別的優化,下面我們講述一些機器學習演算法所用到的技巧。這裡假定用戶都用 SGD 來求解機器學習演算法。TensorFlow 的基本數據流圖模型可以以多種方式用於機器學習應用。我們關心的一個領域是如何加速計算密集型神經網路模型在大型數據集上的訓練。本節描述了我們和其他人為了實現這一點而開發的幾種技術,並說明了如何使用 TensorFlow 實現這些不同的方法。

本小節中的方法假設使用隨機梯度下降法(SGD)對模型進行訓練,使用的小批次包括 100~1000 個樣本。

6.1 數據並行訓練

可以通過數據並行的方式來提升 SGD 的效率,比如,假如每次 SGD 的 mini-batch 是 1000 個樣本,我們可以把模型複製成 10 份放到 10 個 GPU 之上,mini-batch 也被切成 10 份,每份就是 100 個樣本。這個 10 個模型副本並行計算自己 100 個樣本的梯度,然後把這些梯度進行同步規約,最後用於模型參數更新。這就像我們使用 1000個元素的批量大小運行順序 SGD 演算法一樣。在這種情況下, TensorFlow 圖擁有原計算圖中執行大多數計算部分的多個副本,由單個客戶端執行緒驅動這個大型圖的整個訓練循環。下圖頂部對此進行了說明。

圖 7 數據並行

這種方法也可以是非同步的,每一個模型副本非同步地將梯度更新應用於模型參數,不需要彼此等待。在此配置中,每個圖副本都有一個客戶端執行緒。上圖下半部分對此進行了說明。

7. Model Parallel Training

模型並行訓練也很容易用 TensorFlow 表示,這樣對於同一批樣本,模型不同部分可以在不同的計算設備上同時計算。下圖顯示了 LSTM 模型的示例,該模型在三個不同的設備上並行。

0x08 Concurrent Steps for Model Computation PipeLine

在同一設備中對模型計算進行流水線處理也是一個常用的提高利用率的方法,這是通過在同一組設備中運行少量的並發步驟來完成的。它有點類似於非同步數據並行,只是流水線並行發生在同一設備內,而不是在不同設備上複製計算圖。在一個單一的步驟中,在所有設備上的計算可能無法在任何時候完全利用全部設備的並行性,而流水線並行允許 “填補間隙”,這可以充分利用空閑的設備資源。

0xFF 參考

TensorFlow 架構

TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems

TensorFlow : a system for large-scale machine learning

TensorFlow 分散式采坑記