2015年9月，Google于TensorFlow開(kāi)源之際，發(fā)布了TensorFlow白皮書(shū)，介紹了TensorFlow的設(shè)計(jì)理念和實(shí)現(xiàn)方式。現(xiàn)在流行的大部分深度學(xué)習(xí)框架，都基于所謂的“數(shù)據(jù)流圖”編程模型（又稱“計(jì)算圖”），為我們今后的編程提供了一種可選的編程范式。在本文中，我將以問(wèn)答的形式解釋關(guān)于TensorFlow的種種疑問(wèn)。
Qustion 1：“數(shù)據(jù)流圖”編程模型與傳統(tǒng)編程模型的區(qū)別？
首先，“數(shù)據(jù)流圖”的核心是一個(gè)有向圖，圖中的節(jié)點(diǎn)表示運(yùn)算操作，邊表示數(shù)據(jù)，整個(gè)圖展現(xiàn)了數(shù)據(jù)的流動(dòng)，因此稱為數(shù)據(jù)流圖。下圖是一個(gè)示例。
在傳統(tǒng)編程中，雖然我們也是對(duì)數(shù)據(jù)進(jìn)行操作，但基本的三種控制邏輯“順序、選擇、循環(huán)”導(dǎo)致我們只能按照單一的流程處理數(shù)據(jù)，相當(dāng)于數(shù)據(jù)流圖只是一條線，而不是真正的圖。換句話說(shuō)，傳統(tǒng)編程模型解決的是順序操作流程，而數(shù)據(jù)流圖則提供了并行計(jì)算的解決方案。其次，數(shù)據(jù)流圖是數(shù)據(jù)驅(qū)動(dòng)的，而不是指令驅(qū)動(dòng)的。程序只規(guī)定數(shù)據(jù)的流向，而不能規(guī)定每一個(gè)操作何時(shí)執(zhí)行，這就在另一個(gè)層面上提高了并行計(jì)算能力。最后，數(shù)據(jù)流圖（我們這里只探討靜態(tài)數(shù)據(jù)流圖）的定義和執(zhí)行是分開(kāi)的，用戶不能像往常一樣在某個(gè)操作處打斷點(diǎn)查看輸出內(nèi)容，這削弱了該模型的調(diào)試能力，是為性能優(yōu)化而付出的代價(jià)。
Qustion 2：TensorFlow的系統(tǒng)架構(gòu)什么樣？
TensorFlow框架包含三個(gè)模塊，分別是client、master和worker，它們之間的邏輯關(guān)系如下圖所示。
左側(cè)為單機(jī)模式，右側(cè)為分布式模式。client提供用戶使用的編程接口，比如Python、C++ API等。master負(fù)責(zé)接受client的請(qǐng)求，構(gòu)造數(shù)據(jù)流圖，并分配任務(wù)給worker。在單機(jī)模式中，只有一個(gè)worker，負(fù)責(zé)計(jì)算數(shù)據(jù)流圖中的所有計(jì)算任務(wù)。在分布式模式中，master需要為不同的worker分配不同的任務(wù)，以使總用時(shí)最小。
Question 3：分布式模式中，master如何合理地為多個(gè)worker分配任務(wù)？
master需要預(yù)先估算數(shù)據(jù)流圖中各個(gè)節(jié)點(diǎn)的數(shù)據(jù)量和計(jì)算時(shí)長(zhǎng)，估計(jì)數(shù)據(jù)量是為了保證各個(gè)設(shè)備的內(nèi)存占用相差不大，估計(jì)計(jì)算時(shí)長(zhǎng)是為了使總用時(shí)最短。這一步的算法在TensorFlow中稱為Node Placement，即為每個(gè)節(jié)點(diǎn)尋找一個(gè)放置的位置。該算法會(huì)根據(jù)輸入輸出數(shù)據(jù)的規(guī)模、運(yùn)算符種類來(lái)估算。此外，也可以根據(jù)實(shí)際運(yùn)行過(guò)程中的實(shí)測(cè)結(jié)果來(lái)決定。
Question 4：分布式設(shè)備間的數(shù)據(jù)傳輸怎么實(shí)現(xiàn)？
當(dāng)master為各個(gè)設(shè)備劃分好任務(wù)后，這些設(shè)備間不可避免地要進(jìn)行數(shù)據(jù)傳輸。TensorFlow的做法如下圖所示。
首先，對(duì)跨機(jī)器的數(shù)據(jù)傳輸做了一層隔離，A設(shè)備中增加Send節(jié)點(diǎn)，用來(lái)對(duì)外發(fā)送數(shù)據(jù)，B設(shè)備中增加Receive節(jié)點(diǎn)，用來(lái)接收數(shù)據(jù)。這樣，Send和Receive之間的通信與設(shè)備內(nèi)部的通信可以使用完全不同的兩套方案，簡(jiǎn)化了內(nèi)部邏輯。實(shí)際實(shí)現(xiàn)中，Send和Receive之間通過(guò)TCP或RDMA的方式進(jìn)行通信。
Question 5：如何scheduling？
當(dāng)用戶喂入數(shù)據(jù)，啟動(dòng)數(shù)據(jù)流圖后，TensorFlow如何決定各個(gè)節(jié)點(diǎn)的執(zhí)行順序呢？樸素的想法是，每個(gè)節(jié)點(diǎn)執(zhí)行完后通知與之相連的下一個(gè)節(jié)點(diǎn)，但如果下個(gè)節(jié)點(diǎn)有多個(gè)輸入，它仍然無(wú)法啟動(dòng)，必須等到所有輸入都到位后才能開(kāi)始。因此TensorFlow使用了依賴計(jì)數(shù)的機(jī)制，每個(gè)節(jié)點(diǎn)記錄其尚未滿足的依賴的個(gè)數(shù)，當(dāng)個(gè)數(shù)降為0時(shí)，啟動(dòng)該節(jié)點(diǎn)。這一方案完美體現(xiàn)了數(shù)據(jù)驅(qū)動(dòng)的設(shè)計(jì)理念。
Question 6：反向傳播的數(shù)據(jù)依賴怎么處理？
一旦考慮反向傳播，數(shù)據(jù)流圖就沒(méi)那么簡(jiǎn)單了。雖然用戶不必手動(dòng)構(gòu)造反向數(shù)據(jù)流，但TensorFlow會(huì)自動(dòng)為我們構(gòu)造，就像下圖這個(gè)樣子。
仔細(xì)觀察上圖，可以發(fā)現(xiàn)很多跨越數(shù)層的依賴，比如從MatMul到dAdd的灰色連線。這會(huì)導(dǎo)致GPU的內(nèi)存占用急劇增大，因?yàn)閹缀趺總€(gè)階段計(jì)算的中間結(jié)果都不能丟棄。TensorFlow針對(duì)這種情況采取了若干種優(yōu)化措施：(1)用更復(fù)雜的啟發(fā)式算法調(diào)整圖的執(zhí)行順序；(2)反向傳播時(shí)重新計(jì)算前向傳播的中間結(jié)果；(3)把中間結(jié)果保存到CPU內(nèi)存中。
結(jié)語(yǔ)
Tensorflow開(kāi)源至今，已經(jīng)有了長(zhǎng)足的發(fā)展。最新版本的實(shí)現(xiàn)可能與本文所述不再一致，但設(shè)計(jì)者的初衷不變，都是為了提供一個(gè)高擴(kuò)展性、快速、高效的機(jī)器學(xué)習(xí)計(jì)算平臺(tái)。
參考資料
TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems Google Research
