Tensorflow中的控制流和優(yōu)化器指的是什么

只要對tensorflow有一點(diǎn)了解，都應(yīng)該知道graph是tensorflow最基本的一個結(jié)構(gòu)。Tensorflow的所有計(jì)算都是以圖作為依據(jù)的。圖的nodes表示一些基本的數(shù)學(xué)運(yùn)算，比如加法，卷積，pool等。Node使用protoBuf來進(jìn)行描述，包括node的名字，op，input等，詳細(xì)可以參看tensorflow中的node_def.proto文件。Node對應(yīng)的op使用C++來進(jìn)行實(shí)現(xiàn)。圖中的邊表示了數(shù)據(jù)流動的方向以及節(jié)點(diǎn)之間的依賴關(guān)系。比如A->B就表示B必須在A執(zhí)行完之后才能夠執(zhí)行。以下是inception網(wǎng)絡(luò)的圖結(jié)構(gòu)。

當(dāng)了解了tensorflow的一些基本op之后，我們會存在這樣的疑問。對于需要分支跳轉(zhuǎn)，循環(huán)的部分，tensorflow是如何實(shí)現(xiàn)的。比如tf.cond，tf.while_loop這些語句在底層是如何表示的呢？tensorflow定義一些基本的控制原語，通過一定的組合可以完成高層次控制語言的實(shí)現(xiàn)，比如a=op?C:D這樣的語句。

tensorflow控制流的設(shè)計(jì)原則是通過引入最少的控制模塊，利用這些控制模塊可以表示很多復(fù)雜應(yīng)用廣泛的控制過程。這些控制模塊還能夠適應(yīng)并發(fā)，分布式運(yùn)算，同時能夠?qū)崿F(xiàn)自動微分。在tensorflow,一個計(jì)算節(jié)點(diǎn)在執(zhí)行幀（execution frame，類比進(jìn)程的棧幀）里執(zhí)行。控制流原語負(fù)責(zé)創(chuàng)建和管理執(zhí)行。直觀地理解，TF運(yùn)行時建立一個個執(zhí)行幀，在執(zhí)行幀里執(zhí)行所有屬于這個執(zhí)行幀的計(jì)算節(jié)點(diǎn)。執(zhí)行幀可以嵌套（父子關(guān)系）。來自不同執(zhí)行幀且沒有依賴關(guān)系的計(jì)算節(jié)點(diǎn)可以并行計(jì)算。這里介紹5種最基本的控制原語。

1 switch

依據(jù)控制條件p，選擇性將輸入數(shù)據(jù)d傳播到兩個輸出端。

2 merge

Merge算子將一個可用輸入傳給輸出，只要有任意一個輸入可用，switch就可以執(zhí)行。

3 enter

Enter算子依據(jù)執(zhí)行幀唯一標(biāo)識名稱將輸入傳遞到相應(yīng)執(zhí)行幀。Enter算子用于將一個tensor從一個執(zhí)行幀傳遞到子執(zhí)行幀。

4 exit

Exit算子用于將子執(zhí)行幀的數(shù)據(jù)傳遞父執(zhí)行幀。

5 nextIteration

netIteration算子可以將其輸入傳遞到當(dāng)前執(zhí)行幀的下一個iteration。Tensorflow的runtime可以隨時跟蹤執(zhí)行幀中的iteration。任何一個op都有一個唯一的iteration ID進(jìn)行標(biāo)識。

現(xiàn)在我們來看這幾種原子指令是如何實(shí)現(xiàn)條件判斷和循環(huán)的。

Tensorflow中條件判斷cond(pre, fn1, fn2)實(shí)現(xiàn)的偽代碼如下：

首先創(chuàng)建一個條件控制context，這個context會調(diào)用兩個不同的計(jì)算圖。使用哪個計(jì)算圖由條件pre來決定。最后將調(diào)用兩個計(jì)算圖的結(jié)果通過merge節(jié)點(diǎn)輸出到下一個計(jì)算圖。使用merge節(jié)點(diǎn)是為了保證只要有一個圖有了結(jié)果就可以馬上輸送到下一個節(jié)點(diǎn)進(jìn)行后續(xù)計(jì)算。用圖描述如下：

對于循環(huán)語句，tensorflow中使用一下偽代碼來完成：

首先創(chuàng)建一個循環(huán)控制context。然后創(chuàng)建一個enter和merge節(jié)點(diǎn)來導(dǎo)入循環(huán)體變量。使用enter節(jié)點(diǎn)是通過幀名識別這個循環(huán)體從而去執(zhí)行。Merge是將循環(huán)變量傳遞給判斷條件圖，進(jìn)行循環(huán)判定。加入的switch節(jié)點(diǎn)用于對循環(huán)條件判斷的結(jié)果進(jìn)行計(jì)算圖選擇。循環(huán)體內(nèi)部計(jì)算結(jié)果需要進(jìn)行多次循環(huán)，所以進(jìn)入了nextIteration節(jié)點(diǎn)。Switch的false輸出用于終止循環(huán)，所以進(jìn)入exit節(jié)點(diǎn)將最終結(jié)果輸出。

有了這些控制節(jié)點(diǎn)，tensorflow就可以將一個圖分割成多個子圖，并部署到多個硬件執(zhí)行設(shè)備上。在兩個子圖分割處，添加send和receive節(jié)點(diǎn)用于不同設(shè)備之間數(shù)據(jù)通信。Tensorflow對節(jié)點(diǎn)如何分配沒有限制，只要這個節(jié)點(diǎn)可以在這個設(shè)備上執(zhí)行，就可以分配。如果沒有這些控制節(jié)點(diǎn)，那么一幅圖中的一個節(jié)點(diǎn)就只能執(zhí)行一次，有了這些控制節(jié)點(diǎn)，計(jì)算圖就能夠有更多計(jì)算方式。一個節(jié)點(diǎn)可以循環(huán)執(zhí)行多次，還可以被分配到不同設(shè)備執(zhí)行。

Tensorflow可以支持自動微分。當(dāng)用戶建立了計(jì)算圖和定義了loss函數(shù)后，tensorflow會根據(jù)計(jì)算圖的結(jié)構(gòu)建立反向傳播圖。給定一個計(jì)算節(jié)點(diǎn)，可以通過映射到計(jì)算公式方式進(jìn)而求取微分。從而能夠找出其反向傳播的節(jié)點(diǎn)的表示。對于控制節(jié)點(diǎn)來說，enter的反向傳播節(jié)點(diǎn)是exit，switch的反向傳播節(jié)點(diǎn)是merge(對于cond來說)，或者nextIteration+merge(對于while_loop來說)。Merge的反向傳播節(jié)點(diǎn)是switch。nextIteration的反向傳播節(jié)點(diǎn)是identity。Enter的反向傳播節(jié)點(diǎn)是exit。有了這些對應(yīng)關(guān)系，就可以自動來推斷反向傳播圖。從而求取梯度了。而且可以在多個設(shè)備上進(jìn)行計(jì)算分配。

比如對于cond條件判斷，如果其不是loop中的條件判斷，那么其正向傳播圖和反向傳播圖的映射關(guān)系為：

優(yōu)化器是在原始計(jì)算圖基礎(chǔ)上進(jìn)行優(yōu)化，提高計(jì)算在硬件上的效率。優(yōu)化主要有幾個目標(biāo)：簡化圖結(jié)構(gòu)，降低最大的硬件存儲使用率，進(jìn)行硬件友好的圖轉(zhuǎn)化。圖優(yōu)化方法有很多，有些和硬件無關(guān)，有些和硬件的具體實(shí)現(xiàn)細(xì)節(jié)相關(guān)。高層次優(yōu)化是對圖進(jìn)行一定簡化，它對硬件是透明的。通過簡化可以去除一些冗余計(jì)算。比如常數(shù)折疊，多余控制節(jié)點(diǎn)去除等。還有一些利用結(jié)合律，分配律等對公式進(jìn)行簡化，比如：

1)      圖的簡化可以刪除一些冗余計(jì)算，將圖用最終等效結(jié)果替換。比如一個建立tensor的過程：

將tensor的shape創(chuàng)建和數(shù)據(jù)創(chuàng)建合并，直接用常數(shù)替換。這樣就去除了shape創(chuàng)建過程。

2)      常數(shù)折疊可以將兩個以上常數(shù)用一個常數(shù)替代，需要優(yōu)化器進(jìn)行一些計(jì)算。比如：

3）    代數(shù)優(yōu)化利用算術(shù)的性質(zhì)進(jìn)行一定轉(zhuǎn)化。比如：

addN相當(dāng)于硬件上可支持的一個并行計(jì)算單元，可以一次計(jì)算多個輸入。所以可以將連續(xù)的三個加法用一個并行加法替換。

第二個利用了算術(shù)的分配律和結(jié)合律將三個具有相同乘數(shù)提取出來。最后一個對邏輯進(jìn)行了等效轉(zhuǎn)化，從而減少了計(jì)算節(jié)點(diǎn)。

這個matrix+scalar的時候需要對scalar先進(jìn)行廣播，然后再加。轉(zhuǎn)化后減少了廣播次數(shù)。

這兩個消除了冗余計(jì)算。

4）    op融合將多個計(jì)算節(jié)點(diǎn)融合為一個節(jié)點(diǎn)來計(jì)算。這個是和硬件有關(guān)的，比如一個硬件計(jì)算單元可以完成conv+batch_norm，那么就可以實(shí)現(xiàn)這樣的計(jì)算融合，就不需要單獨(dú)多出來一個計(jì)算單元。常用的op融合有：

5）    存儲優(yōu)化的目的是為了降低對片外的訪問頻率，這樣能夠提高數(shù)據(jù)運(yùn)算效率，減少等待數(shù)據(jù)加載時間。