[翻譯] TensorFlow 分佈式之論文篇 "Implementation of Control Flow in TensorFlow"

[翻譯] TensorFlow 分佈式之論文篇 "Implementation of Control Flow in TensorFlow"

讀論文有一種原則是：本領域最經典的論文，近5年最熱的論文，近1年最新的論文。按照這個原則，本文主要介紹一篇Tensorflow 經典論文 Implementation of Control Flow in TensorFlow。

1. 概覽

本文介紹了 TensorFlow 中控制流操作符的當前設計和實現。這是一份基於原始設計的描述性文檔，具體細節請參見實際源代碼。本文內容是：

介紹五個 TensorFlow 的核心操作符，它們是專門爲處理控制流而添加的。
展示高層控制流結構如何基於這五個基礎操作符被編譯進數據流圖。
解釋這些數據流圖如何由 TensorFlow runtime 執行，包括在一組混合設備（如CPU、GPU和TPU）上的分佈式執行方式。
描述如何對控制流結構進行自動求導。

本文圖均來自原始論文。

2. 控制流原語

TensorFlow 中控制流的基本設計原則是：引入一個包含少量操作的簡單原子操作集，在這些操作符之上來表達TensorFlow 應用的複雜控制流。我們希望這些基元是靈活且富有表現力的，可以作爲高級領域特定語言（DSL）的一個良好的編譯目標。它們應該與 TensorFlow 的數據流模型相兼容，並且可以方便實施並行，分佈式執行以及自動微分。如下圖所示，原子操作集之中有五個控制流原語運算符，其中 Switch 和 Merge 組合起來可以實現條件控制。所有五個基元一起組合則可以實現 while 循環。

圖 1 基元

在 TensorFlow 中，每個 op 都在一個執行幀（execution frame）中執行，控制流原語負責創建和管理這些執行幀。對於每個 while 循環，TensorFlow 運行時會設置一個執行幀，並在執行幀內運行 while 循環的所有操作。執行幀可以嵌套。嵌套的 while 循環在嵌套的執行幀中運行。只要執行幀之間沒有數據依賴關係，則來自不同執行幀的操作可以並行運行。

Switch：Switch 運算符會根據輸入控制張量 p 的布爾值，將輸入張量 d 轉發到兩個輸入中的一個。只有兩個輸入都準備好之後，Switch 操作纔會執行。

Merge：Merge 運算符將其可用的輸入之一轉發到其輸出。只要它的任何一個輸入可用，merge 運算符就會執行。如果有多個可用的輸入，則無法確定它的輸出。

Enter(name)：Enter 操作符將其輸入轉發到由給定名稱唯一標識的執行幀。這個 Enter 操作用於將一個執行幀中的張量傳遞給一個子執行幀。對於同一個子執行幀可以有多個 Enter 操作，每個操作都會使子執行幀中的張量可用（異步）。當輸入可用時，Enter 操作將執行。一個新的執行幀在執行該幀第一個 Enter 操作時候被實例化。

Exit：Exit 操作符將一個張量從一個執行幀返回給它的父執行幀。一個執行幀可以有多個 Exit 操作返回到父執行幀，每個操作都異步地將張量傳回給父幀。當一個 Exit 的輸入可用時，該 Exit 操作就被啓用。

NextIteration: 一個 NextIteration 操作符將其輸入轉發到當前執行幀的下一個迭代。TensorFlow 運行時會跟蹤維護執行幀中的迭代信息。一個執行幀中執行的任何操作都有一個唯一的迭代 ID，這使得我們能夠唯一地識別迭代計算中同一操作的不同調用（比如 hile 操作之中，某一個 op 可能會多次執行）。請注意，一個執行幀中可以有多個 NextIteration操作。當執行幀的第 N 次迭代的第一個 NextIteration 操作開始執行時，TensorFlow 運行時就開始進行第 N+1 次迭代。隨着更多的張量通過執行 NextIteration 操作進入下一個迭代，新迭代中更多操作就開始執行。當一個 NextIteration 的輸入可用時，它就被啓用。

3. 控制流結構的編譯

因爲增加了這 5 個控制原語，例如 cond 和 while_loop 這樣的高級編程結構就可以被編譯成數據流圖，從而可以被 TensorFlow 執行。我們接下來看看條件表達式和 while 循環如何在 Tensorflow 內部實現。

3.1 條件表達式

下面是構建條件表達式 cond(pred, fn1, fn2) 數據流圖的高級僞代碼。爲了簡單起見，我們忽略了實際實現中的許多重細節。讀者可以在 control_flow_ops.py 中找到相關的實現細節。

# Build the graph for the true branch 
context_t = CondContext(pred, branch=1) 
res_t = context_t.Call(fn1)

# Build the graph for the false branch 
context_f = CondContext(pred, branch=0) 
res_f = context_f.Call(fn2)

# Add the Merge nodes for the outputs
merges = [Merge([f, t]) for (f, t) in zip(res_f, res_t)] 
return merges

對於條件表達式的每一個分支，我們都會爲條件語境創建一個新的控制流上下文，並在上下文中調用其計算圖構造函數（fn1或fn2）。條件上下文允許我們捕獲任何外部張量（不是在上下文中創建的），並插入一個適當的Switch 操作來確保其進入一個分支。這保證了分支中的任何操作只有在該分支被選擇時纔會執行。由於 TensorFlow 模型的異步執行特點，這些外部張量可能在非常不同的時間變得可用，所以我們爲每個外部張量使用一個 Switch op 來最大化並行度。

因爲每個分支返回一個張量列表（ref_t或res_f），所以我們需要添加一個 Merge 操作來對該結果列表每個輸出的真值/假值進行合併。同樣，輸出可能在不同的時間被計算，所以我們對每個輸出使用一個 Merge 操作，這使我們能夠儘快啓用下游的計算。讓我們來看一個簡單的例子：

圖 2 條件表達式

tf.cond(x<y, lambda: tf.add(x,z), lambda: tf.square(y))

在生成的數據流圖中，Switch 操作被用來控制張量 x、y和z 的流動。在 true/false 分支中，只使用 Switch 操作的真/假輸出。由於 add 的輸入來自 Switch 操作的 true 分支輸出，所以 add 操作只在 x<y 爲真時執行。同樣地，Square 操作只在 x<y 爲假時執行。Add 或 Square 的結果由最後的 Merge 操作發出。如果條件表達式有多個輸出，就會有多個 Merge 操作，每個輸出都有一個 Merge 操作結果。

有很多種使用 Switch 和 Merge 對 cond 進行編碼的方法，我們選擇目前的編碼方式主要是因爲它使 cond 自動求導變得更簡單。

3.2 while 循環

以下是構建 while 循環數據流圖的高層僞代碼：

while_context = WhileContext()
while_context.Enter()

# Add the Enter nodes for each loop variable.
enter_vars = [Enter(x, frame_name) for x in loop_vars]

# Add the Merge nodes. Note that input[1] will be updated later.
merge_vars = [Merge([x,x]) for x in enter_vars]

# Build the loop pred subgraph.
pred_result = pred(*merge_vars)

# Add the Switch nodes.
switch_vars = [Switch(x, pred_result) for x in merge_vars]

# Build the loop body subgraph.
body_result = body(*[x[1] for x in switch_vars])

# Add the NextIteration nodes.
next_vars = [NextIteration(x) for x in body_result]

# Form the cycles for the loop.
for m,v in zip(merge_vars, next_vars):
    m.op._update_input(1,v)

# Add the Exit nodes.
exit_vars = [Exit(x[0]) for x in switch_vars]
while_context.Exit()
return exit_vars

整個 while 循環圖是在 while 循環的控制流上下文之中創建的。這裏的基本思路很簡單。

從循環變量開始，我們爲每個循環變量添加一個 Enter 操作，其後面跟着一個 Merge 操作。然後我們使用其結果（merge_vars）來建立 pred 子圖，pred 子圖將計算循環的終止條件。

在加入 Switch 操作後，我們使用 Switch 的 true 分支輸出來構建 while 循環主體的子圖。循環主體的結果需要進入下一個迭代，所以我們添加 NextIteration 操作，並將其輸出連接到 Merge 操作的第二個輸入。這就形成了循環，這使我們在執行圖的時候可以多次重複運行同一個操作。

Switch 操作的假值輸出是整個 while 循環的輸出，所以我們在假值輸出後面插入了 Exit 操作，並返回 Exit 操作的輸出。與 cond 類似，while 循環的上下文被用來跟蹤 pred 和 body lambdas 中使用的外部張量。這些外部張量被視爲循環常量，我們爲每個這樣的外部張量自動插入一個 Enter 操作，使其可以在 while 循環上下文中訪問。嵌套循環需要添加嵌套的 Enter 操作。

同樣，讓我們看看一個簡單程序的生成圖例子。

圖 3 while 循環

tf.while_loop(lambda i:i<10, lambda i: tf.add(i,1),[0])

在這個例子中，我們只有一個循環變量。如果有多個循環變量，我們需要添加多個 Enter、Merge、Switch、NextIteration 和 Exit 操作。這樣就可以並行執行跨循環和循環內跨迭代的操作。我們省略了在 while 循環中如何處理常量的方法。如果你想了解其細節，請看具體代碼。

cond 和 while_loop 的這種轉換方法可以支持條件表達式和循環的任意嵌套。例如，一個循環體可以調用另一個 while_loop，它將被遞歸地翻譯成一個嵌套的子圖。該翻譯確保每個循環被靜態地分配一個唯一的框架名稱。

4. 實現

TensorFlow 運行時負責數據流圖的執行。讓我們先快速瀏覽一下。爲了在多個設備上運行，TensorFlow 會自動將操作分配到設備集上。TensorFlow 基於設備的具體放置來自動將數據流圖分割成一組子圖，每個設備一個子圖。當一條邊被分區切分時，我們會自動插入一對發送和接收節點，用於在設備間傳輸張量。一對 send 和 recv 使用一個唯一的 key 進行通信，recv 會主動從 send 中提取數據（這裏是特色）。例如，下圖是將一個圖劃分到兩個設備上的結果，TensorFlow 對分區沒有施加任何限制。只要某個節點的計算可以在一個設備上完成，它就可以被分配到該設備上。

圖 4 劃分後的計算圖

當一個子圖被分配到某一個設備之後，這個子圖就被該設備的本地執行器管理。執行器從源節點開始，依次執行準備好的節點。除了合併節點外，一個節點在其所有輸入都可用時，就成爲就緒節點。注意，子圖中的所有 recv 節點都被認爲是源節點。

如果沒有控制流，圖的執行就非常直接。每個節點都僅僅被執行一次，當所有節點都被執行過之後，執行就結束了。控制流引入了相當的複雜性。一個節點現在可以被執行任何次數（包括 0 在內）。執行器需要能夠管理同一節點內多個實例的執行（可能是併發的），並確定圖執行何時會完成。

爲了跟蹤執行過程中產生的張量，我們使用一個元組 d = (value, is_dead, tag) 來標示執行器中的張量，其中 value 是實際的張量，is_dead 是一個布爾值（用來表示該張量是否在一個未執行的條件分支上），而 tag 是唯一標識該張量（以及產生該張量的節點的執行實例）的字符串。直觀地說，tag 定義了一個執行環境，在一個執行環境中，一個節點最多執行一次。標籤是發送/轉發之間通信 key 的一部分，以區分同一發送/轉發節點之間的多個調用。執行者遵循以下執行規則（注意：一個節點的所有輸入必須有相同的標籤。）

Switch(p,d) = (r1,r2)
r1 = (value(d), p || is_dead(d),tag(d))
r2 = (value(d), !p || is_dead(d),tag(d))

Merge(d1,d2) = r
r = if is_dead(d1) then d2 else d1

Enter(d, frame_name) = r
value(r) = value(d)
is_dead(r) = is_dead(d)
tag(r) = tag(d)/frame_name/0

Exit(d) = r
value(r) = value(d)
is_dead(r) = is_dead(d)
tag(r) = tag1 where tag(d)=tag1/frame_name/n

NextIteration(d) = d1
value(d1) = value(d)
is_dead(d1) = is_dead(d)
tag(d1) = tag1/frame_name/(n+1) where tag(d) = tag1/frame_name/n

Op(d1,...,dm) = (r1,...,rn)
value(ri) = Op.Compute(value(d1),...,value(dm)) if !is_dead(ri)
is_dead(ri) = any(is_dead(d1),...,is_dead(dm)), for all i
tag(ri) = tag(d1), for all i

最後一條規則是針對所有非控制流節點的。請注意，只有當所有的輸入都有效時，纔會進行實際的計算。如果有一個無效輸入，我們將跳過計算並向下遊傳播一個 dead 信號。這種 dead 信號的傳播可以被用來支持控制流的分佈式執行。

5. 分佈式條件表達式

對於分佈式執行來說，一個條件表達式可能被切分到多個設備上，如下圖所示：

圖 5 切分表達式

由於任何 recv 節點都是一個隨時無條件啓動的源節點，所以，即使設備 B 上的 recv 節點是在條件表達式的未選擇分支之內，它也可能會執行。爲了使未選擇分支上的 recv 的執行合理化，我們在設備間把 is_dead 標誌通過 send 節點發送到 recv 節點。傳播可以在任何數量的設備上繼續進行。這個簡單的傳播機制可以處理嵌套條件的分佈式執行，也有助於 while 循環的分佈式執行。

6. 分佈式的 while 循環

對於分佈式執行，一個 while 循環，特別是循環主體，可以被切分到多個設備上。如果我們簡單地應用切分方案：只是爲跨設備的邊插入 send/recv 節點，那麼設備上的本地執行器將缺少足夠的信息來正確運行 while 循環。

圖 6 切分控制流簡單方案

讓我們用一個簡單的例子來說明這些問題。在上面的例子中，Op 在循環體中，被分配給設備B。一個簡單切分會將 Switch 到 Op 的邊拆分，插入一對 send/recv 節點，由這對節點完成跨設備數據傳輸。然而，這是不可行的，因爲設備 B 不知道 recv 和 Op 節點是一個 while 循環的一部分，這樣設備 B 在一個迭代後就會終止執行。解決方案是重寫數據流圖，在每個分區添加一個控制循環狀態機（如下圖設備 B 的右下角所示）。控制循環 Enter 節點是一個標量 0。

圖 7 切分控制流改進方案

這些控制循環提供了足夠的信息，這樣通過發送/接收節點相互通信，就可以使設備上的執行器能夠像以前一樣獨立運行。請注意，圖中的虛線是控制邊。讓我們先看一下基本用例，即 while 循環只運行 0 次迭代。

在設備 A 上，節點 Enter、Merge、P 和 Switc 依次被執行。因爲 P 是 false，所以連接到 Switch 的 Send 會向設備 B 傳播一個死信號，這樣 Exit 也會運行，從而使循環之外依賴這個 Exit 的節點能夠同時執行。連接到P 的 Send將向設備 B 發送布爾張量 False，這樣 Recv 也可以被執行，其會等待來自設備 B 的值。
在設備 B 上，Enter 觸發了循環，接下來依次執行節點 Enter 和 Merge。Merge 的執行使兩個 Recv 得以執行。Switch 的 Recv 會收到 False，所以 Next 會得到一個死張量，於是停止了循環。Op 的 Recv 會得到一個死張量，所以 Op 的 Send 會把一個死張量送回設備 A，此時，設備 B 沒有未完成的操作，所以執行結束。
在設備 A 上，Recv for Next 得到了一個死張量。Next 運行，由於它停止了死循環的傳播，設備 A 沒有未完成的操作，所以執行結束。

我們接下來看看 while 循環運行一個或多個迭代。

在設備 A 上，由於 P 在第一次迭代時爲真，一個實數張量被髮送到設備 B。同時 Recv 被執行，等待來自設備B 返回的值。
在設備 B 上，控制循環狀態機運行並啓用 Recv。Recv 爲 Op 從設備 A 得到一個實數張量；Op 被執行，Send 將一個實數張量送回設備 A。執行 Next 和 Merge，進一步啓用下一個迭代的 Recv。
在設備 A 上，Recv 得到一個實數張量。然後執行 Next、Merge 和 P。根據 P 的值，將執行基本情況或新的迭代。

請注意，在執行過程中存在大量的並行性。例如，設備 B 一旦收到 P 的值，就可以開始下一個迭代或退出。一個參與設備可以有多個迭代在並行運行，而且兩個參與設備可以同時在同一個循環的不同迭代中工作。

分佈式執行 while 循環的開銷是每個參與設備在每次迭代時都需要從產生 P 的設備那裏接收一個布爾張量，考慮到執行中的並行性，開銷在很大程度上應該是與計算重疊，因此可以忽略。

下面顯示了當一個 while 循環被劃分到多個設備上時，數據流圖是什麼樣子的。一個控制循環被添加到每個分區中，並控制 while 循環中的 Recvs。重寫後的圖在語義上與原始圖是等價的。

圖 8 重寫的計算圖

對於嵌套的 while 循環，我們按如下方式把控制循環堆疊起來。注意，如果一個設備只有外層循環的節點，我們將不會在其上添加任何與內層循環有關的控制循環結構。

圖 9 嵌套

7. 自動微分

TensorFlow 支持自動求導。例如，用戶可以定義一個帶有損失函數的神經網絡，而 TensorFlow 將自動推導並構建反向傳播數據流圖。本節解釋了 TensorFlow 如何在有 cond 和 while_loop 的情況下自動構建反向傳播圖。我們假設讀者對自動反向傳播的工作方式有一定的瞭解。(參見鏈接 [1]，這是一篇關於反向傳播的優秀文章）。

反向傳播算法以反向順序遍歷前向圖中的操作，並通過調用操作註冊的梯度函數逐步構建梯度圖。一個操作的梯度函數定義了計算該操作梯度的子圖。梯度函數可能會使用到運算的輸入/輸出值，因此在前向計算中產生的一些張量將被保留一段時間，直到它在反向傳播之中被使用。例如，下面顯示了一個前向運算和它的梯度圖。G(Op) 是Op 的梯度子圖。x 和 y 的值將被保存在內存中，直到 G(Op) 被執行。

圖 10 反向傳播

一旦構建了整個數據流圖，TensorFlow 運行時就會自動對圖進行分割，並將執行分佈在多個設備上。因此，TensorFlow 中的梯度計算也將被分配到多個設備上運行。

直觀地講，在 cond 和 while_loop 的上下文之中，控制流算子的反向傳播以如下方式進行反向傳播。Exit 的梯度是 Enter；Switch 的梯度是 Merge（對於cond）或者 NextIteration 之後接着一個 Merge（對於while_loop）；Merge 的梯度是 Switch；NextIteration 的梯度是 Identity；Enter 的梯度是 Exit。TensorFlow 支持嵌套條件和while循環的反向傳播。

7.1 條件表達式的反向傳播

直觀地說，cond(p, fn1, fn2) 的梯度爲 cond(p, g_fn1, g_fn2)，其中 g_fn1 和 g_fn2 分別爲 fn1 和 fn2 的梯度。下面顯示了當 cond 沒有嵌套在 while 循環中，cond 的基本反向傳播操作。我們假設 Op 位於 cond 的 true 分支上。如果 cond 被嵌套在 while 循環，那麼它需要做更多的工作來記住前向循環每次迭代的 p 值。我們將在後面看while 循環的反向傳播時討論這個問題。

圖 10 條件表達式的反向傳播

前向傳播之中的 Merge 在後向傳播之中被轉化爲 Switch，它使用與前向 Switch 相同的謂詞 p。梯度 g 被反推到Switch 的兩個分支。

前向 Switch 被轉化爲 Merge。如果前向 Switch 中只有一個分支在前向傳播之中被用到了，我們會添加一個零輸入到反向傳播的 Merge，如下圖所示，以確保在反向傳播之中總有一個活躍的梯度流經 Merge。這個零輸入被一個 Switch 來控制，所以它只在 p 爲 false 時纔會被髮送到 Merge。

圖 12 Switch 轉換

7.2 While 循環的反向傳播

直觀地說，while_loop(pred, body) 的梯度也是以 while loop 的形式存在。

def pred(i, _): return i < N

while_loop(pred, g_body, [0] + g_vars)

其中 N 是前向傳播 while 循環運行的迭代次數，g_body 是前向循環體的梯度，g_vars 是循環變量的初始值。我們將在後面看到，g_vars 包括前向 while 循環變量的初始梯度。下面是一個 while 循環的前向傳播和反向傳播圖。

圖 13 While 循環的反向傳播

請注意，Backprop 循環由 N 控制，即前向循環運行的迭代次數。這意味着我們假設 pred 是不可訓練的。G(Body) 是 Body 的梯度。Body 可能再次包含 while 循環，所以這個結構可能會遞歸地出現，以處理嵌套的 while 循環。

到目前爲止，這個描述是相當過度簡化了。實際上，在圖的構造過程中，N 並不是靜態已知的。更重要的是，G(Body) 可能會使用前向傳播過程中產生的值，我們希望保留這些值，以避免在反推過程中重新計算它們。解決方案是重寫前向 while 循環的圖，對於反向傳播之中需要的值，增加計算和/或保存的邏輯。

爲了計算 N，我們在前向 while 循環中加入以下子圖（計算 N 的邏輯）。因此，N 將由前向循環動態計算，並作爲後向循環的計數循環變量的初始值。

圖 14 計算邏輯

爲了在反向傳播循環中重用前向傳播計算出來的數值，我們在構建反向傳播 while 循環的過程中，自動檢測反向傳播中需要的前向值。對於每個這樣的前向值 x，我們自動引入一個堆棧，並在前向循環中添加節點，以便在每次迭代時將其值保存到堆棧中。反向傳播循環以相反的順序使用堆棧中的值。堆棧位於前向和反向傳播循環之外，由兩個循環共享（所以下圖有兩個 Enter）。

圖 15 循環共享

實際的計算圖構造實際上比這更微妙和複雜。下面是一些問題。

爲了保證正確性，我們需要確保堆棧的 push 和 pop 是按其各自循環的迭代來排序的。我們還需要確保前向傳播的堆棧必須在後向傳播的堆棧之前完成排序。這些順序是通過控制邊來完成的。
爲了提高性能，我們使堆棧 push 和 pop 操作成爲異步的，因此它們可以與實際計算並行運行。例如，op（甚至是未來的迭代）可以與 push 並行運行。
如果 op 在一個嵌套在 while 循環內的 cond 裏面，那麼入棧和出棧操作必須由 cond 的謂詞進行適當的保護。
如果某個值在反向傳播之中被縮減操作（如 Shape、Rank或Size）處理，我們將縮減操作移到前向循環中以減少內存的使用。

如前所述，Enter 的梯度是 Exit。對於循環變量，這就是它的全部作用。對於循環常量，我們還添加了一個子圖來累積它們的梯度，如下圖所示。

圖 16 累計梯度

假設 x 是前向傳播中的一個循環常數。在 Backprop 中，每次迭代都會爲 x 產生一個 partial gradient。因此，我們在反向傳播過程中添加小的累積子圖，然後將所有這些部分梯度加在一起。最終結果 \(g_x\) 是所有偏導數的總和。注意，積累是 eagerly 地進行的，以並行迭代的次數爲界。這與 static unrolling 不同，在 static unrolling 中，AddN 需要所有的部分梯度在同一時間生效。

這種結構對嵌套條件和循環都有效。對於嵌套在 while 循環中的條件式，我們引入一個堆棧來保存每次前向迭代的謂詞值，並在反向 prop 中使用堆棧中的值（以相反的順序）。對於嵌套的循環，當我們遇到嵌套在循環體中的內部 while 循環時，會遞歸地調用這個結構。

一個重要的優化是內存交換（memory swapping）。正如我們所看到的，對於每個在 backprop 中需要的前向值 v，我們將其在所有迭代中的值 \(v_1,...,v_N\)保存在一個堆棧中，所以我們會在 backprop 中重使它們。這對於在內存有限的設備（如GPU）上進行訓練是一個限制。我們使用內存交換來異步地將存儲在堆棧中的值從 GPU 移動到 CPU，並在 Backprop 中需要時將它們移回 GPU 內存中。