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

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

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

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

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

1. 概覽

本文介紹了 TensorFlow 中控制流操作符的當(dāng)前設(shè)計和實現(xiàn)。這是一份基于原始設(shè)計的描述性文檔，具體細(xì)節(jié)請參見實際源代碼。本文內(nèi)容是：

介紹五個 TensorFlow 的核心操作符，它們是專門為處理控制流而添加的。
展示高層控制流結(jié)構(gòu)如何基于這五個基礎(chǔ)操作符被編譯進(jìn)數(shù)據(jù)流圖。
解釋這些數(shù)據(jù)流圖如何由 TensorFlow runtime 執(zhí)行，包括在一組混合設(shè)備（如CPU、GPU和TPU）上的分布式執(zhí)行方式。
描述如何對控制流結(jié)構(gòu)進(jìn)行自動求導(dǎo)。

本文圖均來自原始論文。

2. 控制流原語

TensorFlow 中控制流的基本設(shè)計原則是：引入一個包含少量操作的簡單原子操作集，在這些操作符之上來表達(dá)TensorFlow 應(yīng)用的復(fù)雜控制流。我們希望這些基元是靈活且富有表現(xiàn)力的，可以作為高級領(lǐng)域特定語言（DSL）的一個良好的編譯目標(biāo)。它們應(yīng)該與 TensorFlow 的數(shù)據(jù)流模型相兼容，并且可以方便實施并行，分布式執(zhí)行以及自動微分。如下圖所示，原子操作集之中有五個控制流原語運算符，其中 Switch 和 Merge 組合起來可以實現(xiàn)條件控制。所有五個基元一起組合則可以實現(xiàn) while 循環(huán)。

圖 1 基元

在 TensorFlow 中，每個 op 都在一個執(zhí)行幀（execution frame）中執(zhí)行，控制流原語負(fù)責(zé)創(chuàng)建和管理這些執(zhí)行幀。對于每個 while 循環(huán)，TensorFlow 運行時會設(shè)置一個執(zhí)行幀，并在執(zhí)行幀內(nèi)運行 while 循環(huán)的所有操作。執(zhí)行幀可以嵌套。嵌套的 while 循環(huán)在嵌套的執(zhí)行幀中運行。只要執(zhí)行幀之間沒有數(shù)據(jù)依賴關(guān)系，則來自不同執(zhí)行幀的操作可以并行運行。

Switch：Switch 運算符會根據(jù)輸入控制張量 p 的布爾值，將輸入張量 d 轉(zhuǎn)發(fā)到兩個輸入中的一個。只有兩個輸入都準(zhǔn)備好之后，Switch 操作才會執(zhí)行。

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

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

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

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

3. 控制流結(jié)構(gòu)的編譯

因為增加了這 5 個控制原語，例如 cond 和 while_loop 這樣的高級編程結(jié)構(gòu)就可以被編譯成數(shù)據(jù)流圖，從而可以被 TensorFlow 執(zhí)行。我們接下來看看條件表達(dá)式和 while 循環(huán)如何在 Tensorflow 內(nèi)部實現(xiàn)。

3.1 條件表達(dá)式

下面是構(gòu)建條件表達(dá)式 cond(pred, fn1, fn2) 數(shù)據(jù)流圖的高級偽代碼。為了簡單起見，我們忽略了實際實現(xiàn)中的許多重細(xì)節(jié)。讀者可以在 control_flow_ops.py 中找到相關(guān)的實現(xiàn)細(xì)節(jié)。

# 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

對于條件表達(dá)式的每一個分支，我們都會為條件語境創(chuàng)建一個新的控制流上下文，并在上下文中調(diào)用其計算圖構(gòu)造函數(shù)（fn1或fn2）。條件上下文允許我們捕獲任何外部張量（不是在上下文中創(chuàng)建的），并插入一個適當(dāng)?shù)腟witch 操作來確保其進(jìn)入一個分支。這保證了分支中的任何操作只有在該分支被選擇時才會執(zhí)行。由于 TensorFlow 模型的異步執(zhí)行特點，這些外部張量可能在非常不同的時間變得可用，所以我們?yōu)槊總€外部張量使用一個 Switch op 來最大化并行度。

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

圖 2 條件表達(dá)式

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

在生成的數(shù)據(jù)流圖中，Switch 操作被用來控制張量 x、y和z 的流動。在 true/false 分支中，只使用 Switch 操作的真/假輸出。由于 add 的輸入來自 Switch 操作的 true 分支輸出，所以 add 操作只在 x<y 為真時執(zhí)行。同樣地，Square 操作只在 x<y 為假時執(zhí)行。Add 或 Square 的結(jié)果由最后的 Merge 操作發(fā)出。如果條件表達(dá)式有多個輸出，就會有多個 Merge 操作，每個輸出都有一個 Merge 操作結(jié)果。

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

3.2 while 循環(huán)

以下是構(gòu)建 while 循環(huán)數(shù)據(jù)流圖的高層偽代碼：

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 循環(huán)圖是在 while 循環(huán)的控制流上下文之中創(chuàng)建的。這里的基本思路很簡單。

從循環(huán)變量開始，我們?yōu)槊總€循環(huán)變量添加一個 Enter 操作，其后面跟著一個 Merge 操作。然后我們使用其結(jié)果（merge_vars）來建立 pred 子圖，pred 子圖將計算循環(huán)的終止條件。

在加入 Switch 操作后，我們使用 Switch 的 true 分支輸出來構(gòu)建 while 循環(huán)主體的子圖。循環(huán)主體的結(jié)果需要進(jìn)入下一個迭代，所以我們添加 NextIteration 操作，并將其輸出連接到 Merge 操作的第二個輸入。這就形成了循環(huán)，這使我們在執(zhí)行圖的時候可以多次重復(fù)運行同一個操作。

Switch 操作的假值輸出是整個 while 循環(huán)的輸出，所以我們在假值輸出后面插入了 Exit 操作，并返回 Exit 操作的輸出。與 cond 類似，while 循環(huán)的上下文被用來跟蹤 pred 和 body lambdas 中使用的外部張量。這些外部張量被視為循環(huán)常量，我們?yōu)槊總€這樣的外部張量自動插入一個 Enter 操作，使其可以在 while 循環(huán)上下文中訪問。嵌套循環(huán)需要添加嵌套的 Enter 操作。

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

圖 3 while 循環(huán)

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

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

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

4. 實現(xiàn)

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

圖 4 劃分后的計算圖

當(dāng)一個子圖被分配到某一個設(shè)備之后，這個子圖就被該設(shè)備的本地執(zhí)行器管理。執(zhí)行器從源節(jié)點開始，依次執(zhí)行準(zhǔn)備好的節(jié)點。除了合并節(jié)點外，一個節(jié)點在其所有輸入都可用時，就成為就緒節(jié)點。注意，子圖中的所有 recv 節(jié)點都被認(rèn)為是源節(jié)點。

如果沒有控制流，圖的執(zhí)行就非常直接。每個節(jié)點都僅僅被執(zhí)行一次，當(dāng)所有節(jié)點都被執(zhí)行過之后，執(zhí)行就結(jié)束了。控制流引入了相當(dāng)?shù)膹?fù)雜性。一個節(jié)點現(xiàn)在可以被執(zhí)行任何次數(shù)（包括 0 在內(nèi)）。執(zhí)行器需要能夠管理同一節(jié)點內(nèi)多個實例的執(zhí)行（可能是并發(fā)的），并確定圖執(zhí)行何時會完成。

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

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

最后一條規(guī)則是針對所有非控制流節(jié)點的。請注意，只有當(dāng)所有的輸入都有效時，才會進(jìn)行實際的計算。如果有一個無效輸入，我們將跳過計算并向下游傳播一個 dead 信號。這種 dead 信號的傳播可以被用來支持控制流的分布式執(zhí)行。

5. 分布式條件表達(dá)式

對于分布式執(zhí)行來說，一個條件表達(dá)式可能被切分到多個設(shè)備上，如下圖所示：

圖 5 切分表達(dá)式

由于任何 recv 節(jié)點都是一個隨時無條件啟動的源節(jié)點，所以，即使設(shè)備 B 上的 recv 節(jié)點是在條件表達(dá)式的未選擇分支之內(nèi)，它也可能會執(zhí)行。為了使未選擇分支上的 recv 的執(zhí)行合理化，我們在設(shè)備間把 is_dead 標(biāo)志通過 send 節(jié)點發(fā)送到 recv 節(jié)點。傳播可以在任何數(shù)量的設(shè)備上繼續(xù)進(jìn)行。這個簡單的傳播機(jī)制可以處理嵌套條件的分布式執(zhí)行，也有助于 while 循環(huán)的分布式執(zhí)行。

6. 分布式的 while 循環(huán)

對于分布式執(zhí)行，一個 while 循環(huán)，特別是循環(huán)主體，可以被切分到多個設(shè)備上。如果我們簡單地應(yīng)用切分方案：只是為跨設(shè)備的邊插入 send/recv 節(jié)點，那么設(shè)備上的本地執(zhí)行器將缺少足夠的信息來正確運行 while 循環(huán)。

圖 6 切分控制流簡單方案

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

圖 7 切分控制流改進(jìn)方案

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

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

我們接下來看看 while 循環(huán)運行一個或多個迭代。

在設(shè)備 A 上，由于 P 在第一次迭代時為真，一個實數(shù)張量被發(fā)送到設(shè)備 B。同時 Recv 被執(zhí)行，等待來自設(shè)備B 返回的值。
在設(shè)備 B 上，控制循環(huán)狀態(tài)機(jī)運行并啟用 Recv。Recv 為 Op 從設(shè)備 A 得到一個實數(shù)張量；Op 被執(zhí)行，Send 將一個實數(shù)張量送回設(shè)備 A。執(zhí)行 Next 和 Merge，進(jìn)一步啟用下一個迭代的 Recv。
在設(shè)備 A 上，Recv 得到一個實數(shù)張量。然后執(zhí)行 Next、Merge 和 P。根據(jù) P 的值，將執(zhí)行基本情況或新的迭代。

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

分布式執(zhí)行 while 循環(huán)的開銷是每個參與設(shè)備在每次迭代時都需要從產(chǎn)生 P 的設(shè)備那里接收一個布爾張量，考慮到執(zhí)行中的并行性，開銷在很大程度上應(yīng)該是與計算重疊，因此可以忽略。

下面顯示了當(dāng)一個 while 循環(huán)被劃分到多個設(shè)備上時，數(shù)據(jù)流圖是什么樣子的。一個控制循環(huán)被添加到每個分區(qū)中，并控制 while 循環(huán)中的 Recvs。重寫后的圖在語義上與原始圖是等價的。

圖 8 重寫的計算圖

對于嵌套的 while 循環(huán)，我們按如下方式把控制循環(huán)堆疊起來。注意，如果一個設(shè)備只有外層循環(huán)的節(jié)點，我們將不會在其上添加任何與內(nèi)層循環(huán)有關(guān)的控制循環(huán)結(jié)構(gòu)。

圖 9 嵌套

7. 自動微分

TensorFlow 支持自動求導(dǎo)。例如，用戶可以定義一個帶有損失函數(shù)的神經(jīng)網(wǎng)絡(luò)，而 TensorFlow 將自動推導(dǎo)并構(gòu)建反向傳播數(shù)據(jù)流圖。本節(jié)解釋了 TensorFlow 如何在有 cond 和 while_loop 的情況下自動構(gòu)建反向傳播圖。我們假設(shè)讀者對自動反向傳播的工作方式有一定的了解。(參見鏈接 [1]，這是一篇關(guān)于反向傳播的優(yōu)秀文章）。

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

圖 10 反向傳播

一旦構(gòu)建了整個數(shù)據(jù)流圖，TensorFlow 運行時就會自動對圖進(jìn)行分割，并將執(zhí)行分布在多個設(shè)備上。因此，TensorFlow 中的梯度計算也將被分配到多個設(shè)備上運行。

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

7.1 條件表達(dá)式的反向傳播

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

圖 10 條件表達(dá)式的反向傳播

前向傳播之中的 Merge 在后向傳播之中被轉(zhuǎn)化為 Switch，它使用與前向 Switch 相同的謂詞 p。梯度 g 被反推到Switch 的兩個分支。

前向 Switch 被轉(zhuǎn)化為 Merge。如果前向 Switch 中只有一個分支在前向傳播之中被用到了，我們會添加一個零輸入到反向傳播的 Merge，如下圖所示，以確保在反向傳播之中總有一個活躍的梯度流經(jīng) Merge。這個零輸入被一個 Switch 來控制，所以它只在 p 為 false 時才會被發(fā)送到 Merge。

圖 12 Switch 轉(zhuǎn)換

7.2 While 循環(huán)的反向傳播

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

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

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

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

圖 13 While 循環(huán)的反向傳播

請注意，Backprop 循環(huán)由 N 控制，即前向循環(huán)運行的迭代次數(shù)。這意味著我們假設(shè) pred 是不可訓(xùn)練的。G(Body) 是 Body 的梯度。Body 可能再次包含 while 循環(huán)，所以這個結(jié)構(gòu)可能會遞歸地出現(xiàn)，以處理嵌套的 while 循環(huán)。

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

為了計算 N，我們在前向 while 循環(huán)中加入以下子圖（計算 N 的邏輯）。因此，N 將由前向循環(huán)動態(tài)計算，并作為后向循環(huán)的計數(shù)循環(huán)變量的初始值。

圖 14 計算邏輯

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

圖 15 循環(huán)共享

實際的計算圖構(gòu)造實際上比這更微妙和復(fù)雜。下面是一些問題。

為了保證正確性，我們需要確保堆棧的 push 和 pop 是按其各自循環(huán)的迭代來排序的。我們還需要確保前向傳播的堆棧必須在后向傳播的堆棧之前完成排序。這些順序是通過控制邊來完成的。
為了提高性能，我們使堆棧 push 和 pop 操作成為異步的，因此它們可以與實際計算并行運行。例如，op（甚至是未來的迭代）可以與 push 并行運行。
如果 op 在一個嵌套在 while 循環(huán)內(nèi)的 cond 里面，那么入棧和出棧操作必須由 cond 的謂詞進(jìn)行適當(dāng)?shù)谋Ｗo(hù)。
如果某個值在反向傳播之中被縮減操作（如 Shape、Rank或Size）處理，我們將縮減操作移到前向循環(huán)中以減少內(nèi)存的使用。

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

圖 16 累計梯度

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

這種結(jié)構(gòu)對嵌套條件和循環(huán)都有效。對于嵌套在 while 循環(huán)中的條件式，我們引入一個堆棧來保存每次前向迭代的謂詞值，并在反向 prop 中使用堆棧中的值（以相反的順序）。對于嵌套的循環(huán)，當(dāng)我們遇到嵌套在循環(huán)體中的內(nèi)部 while 循環(huán)時，會遞歸地調(diào)用這個結(jié)構(gòu)。

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