Tensorflow核心代碼解析之計(jì)算圖篇其一:計(jì)算圖結(jié)構(gòu)初探

介紹

當(dāng)今計(jì)算機(jī)科學(xué)給人對(duì)未來最大想象的莫過于人工智能的大規(guī)模應(yīng)用前景。它對(duì)于人類文明進(jìn)步所帶來的潛在貢獻(xiàn)可以被視為與四大發(fā)明、蒸汽機(jī)、電力、計(jì)算機(jī)等人類史標(biāo)桿性工具具有同等的地位。當(dāng)今人工智能的蓬勃發(fā)展主要是機(jī)器學(xué)習(xí)尤其是深度學(xué)習(xí)的大規(guī)模成功應(yīng)用。憑著日益增多的海量數(shù)據(jù),快速發(fā)展的計(jì)算機(jī)并行計(jì)算能力,快速迭代、高效更新的各種模型、算法、策略以及各個(gè)國(guó)家、政府、大企業(yè)對(duì)它的日益重視與超高的資金、政策投入,當(dāng)今AI的發(fā)展速度真可謂一日千里!

深度學(xué)習(xí)最終要與具體行業(yè)場(chǎng)景有效率地結(jié)合起來才能發(fā)揮出其效益來。當(dāng)下的整個(gè)深度學(xué)習(xí)已經(jīng)行成了良好的產(chǎn)業(yè)鏈體系。最下端的是用于深度學(xué)習(xí)加速的各種硬件芯片如CPU/GPU/FPGA/ASIC專用芯片等。目前此領(lǐng)域里面GPU憑其優(yōu)良的超多弱核心并行計(jì)算能力獨(dú)領(lǐng)風(fēng)騷;但CPU在推理加速、成本等方面也挺有競(jìng)爭(zhēng)力;FPGA憑其靈活性也非常適宜于進(jìn)行AI方案的原型設(shè)計(jì),但因其開發(fā)難度較大,生態(tài)相對(duì)較缺乏,當(dāng)下大公司里面大規(guī)模部署應(yīng)用FPGA的唯有微軟;至于ASIC專用芯片,可謂是給了諸多有心在AI半導(dǎo)體上面實(shí)現(xiàn)彎道超車的公司一個(gè)很好的機(jī)會(huì),尤其是那些有著大規(guī)模機(jī)器學(xué)習(xí)用戶,可以基于上層封裝提供AI云計(jì)算實(shí)例服務(wù)的公司如Google就聲勢(shì)浩大地推出了自己快速迭代著的TPU,號(hào)稱常規(guī)模型(如Resnet-50)加速快于最新的GPU,同時(shí)功耗更為節(jié)省,其它像Amazon,F(xiàn)acebook,Ali等云計(jì)算公司也在搞自己的AI專用ASIC芯片。此外一些創(chuàng)業(yè)型公司也對(duì)此雄心勃勃,國(guó)內(nèi)已經(jīng)涌現(xiàn)出了一堆此類的獨(dú)角獸像賣自家礦機(jī)發(fā)了大財(cái)?shù)腂it大陸,技術(shù)勢(shì)力雄厚的地平線等等。還有些手機(jī)大廠則將精力用在了終端一側(cè)的AI芯片研發(fā)像華為、蘋果等已經(jīng)有了自己的AI芯片并部署在了自己的手機(jī)新品當(dāng)中,其它像小米、三星等也在紛紛跟進(jìn)。需要提下的是ARM這個(gè)在移動(dòng)時(shí)代的主要得意者也于最近發(fā)布了自己家設(shè)計(jì)的用于AI加速的各種硬件IP。

AI芯片

計(jì)算芯片向上走則是將一些基本運(yùn)算如矩陣乘積、各類型卷積運(yùn)算等結(jié)合硬件平臺(tái)優(yōu)化過了的數(shù)學(xué)計(jì)算庫如用于Intel CPU端的MKL/MKLDNN,用于nVidia GPU的CUDA/cuDNN,用于FPGA專用DNN網(wǎng)絡(luò)加速過的openCL,還有針對(duì)各大ASIC芯片產(chǎn)商針對(duì)自己ASIC加速過了的種種計(jì)算庫套件。這些數(shù)學(xué)計(jì)算庫基本由AI芯片產(chǎn)商自己來完成,目的即在于借用軟件的力量給自家的硬件以強(qiáng)大的驅(qū)動(dòng)力。一般他們會(huì)選擇將優(yōu)化過了的核心程序開源、公布出來為自己的客戶所借鑒、使用,最終通過AI芯片的出售來獲得價(jià)值。值得一提的是nVIDIA在基于CUDA/cuDNN上面的多年耕耘,相關(guān)社區(qū)、生態(tài)的耐心培養(yǎng)直接帶來了它們今天的巨大成功。現(xiàn)在半導(dǎo)體公司已經(jīng)再不同往日只需要賣出芯片即可了。芯片相關(guān)的軟件庫(并行計(jì)算庫等)的性能,對(duì)用戶提供API的友好性,用戶社區(qū)的培養(yǎng),用戶支持的力度等等軟實(shí)力真正是愈益重要。

底層計(jì)算庫再向上走則是使用這些優(yōu)化過了的數(shù)學(xué)計(jì)算庫來完成基本類型計(jì)算,然后將之抽象封裝后向上提供出友好用戶API的深度學(xué)習(xí)框架。這些深度學(xué)習(xí)框架是我們軟件人員開發(fā)應(yīng)用的基本工具。當(dāng)前最流行的框架有Google的Tensorflow,F(xiàn)acebook的Pytorch,Amazon的Mxnet,微軟的CNTK,當(dāng)然還有傳統(tǒng)社區(qū)在維護(hù)的bvlc/Caffe等等。它們大都提供類似的功能,相似的API用于用戶程序構(gòu)建計(jì)算圖,并能將圖方便地導(dǎo)入、導(dǎo)出為序列化文件,還提供了基于Framework level對(duì)圖的一些優(yōu)化如合并(fusion),去重(典型的如CSE),并行計(jì)算(通過使用OMP等并行庫)等,此外還有一些功能如用于進(jìn)行內(nèi)存分配、管理及線程執(zhí)行、調(diào)度、檢測(cè)的session/workspace等,當(dāng)然還有用于具體執(zhí)行某計(jì)算的op / kernel等,這也是常規(guī)計(jì)算優(yōu)化的核心所在。

在這些計(jì)算框架中,無疑Google brain團(tuán)隊(duì)開發(fā)的Tensorflow是最為流行的。它的框架設(shè)計(jì)最為復(fù)雜,可以天生地支持模型并行訓(xùn)練、推理等,它的背后有一個(gè)google強(qiáng)大的開發(fā)團(tuán)隊(duì)在快速迭代、開發(fā),它的底下也有集成當(dāng)前最好的像cuDNN/mklDNN/TensorRT等加速技術(shù),它的用戶社區(qū)也已經(jīng)非常完善、活躍(作為程序員這個(gè)還是蠻重要的。畢竟在APP開發(fā)中出了問題,肯定都希望通過在網(wǎng)上翻一下看有沒有人踩過類似的坑以來快速解決問題。)。

當(dāng)下對(duì)于如何使用Tensorflow來開發(fā)一個(gè)AI程序,構(gòu)建深度學(xué)習(xí)模型并進(jìn)行訓(xùn)練或推理的文章已經(jīng)很多了。本系列單元中筆者想試著跟大家一起理一下它框架核心的一些代碼實(shí)現(xiàn)。無益它對(duì)于我們基于TF做一些開發(fā),加深對(duì)TF的理解是很有幫助的。此外TF框架的設(shè)計(jì)、代碼實(shí)現(xiàn)非常良好,對(duì)它們的理解、梳理清楚對(duì)于我們?nèi)粘5能浖O(shè)計(jì)、開發(fā)也會(huì)有較強(qiáng)的工程借鑒意義。

計(jì)算圖

計(jì)算圖(Graph)描述了一組需要依次序完成的計(jì)算單元以及表示這些計(jì)算單元之間相互依賴的關(guān)系。一般的深度學(xué)習(xí)模型都會(huì)被分化組裝成一個(gè)單向無環(huán)圖(DAG)來執(zhí)行。圖當(dāng)中的結(jié)點(diǎn)(node)用來表示某一具體的計(jì)算單元(如Multmul結(jié)點(diǎn)表示兩個(gè)張量之間的乘積,Conv結(jié)點(diǎn)則表示兩個(gè)張量之間的卷積計(jì)算)。圖上的片(edge)則被用來表示兩個(gè)結(jié)點(diǎn)之間的依賴關(guān)系。比如A結(jié)點(diǎn)的第i個(gè)輸出來自于B結(jié)點(diǎn)的第j輸入,那么就會(huì)構(gòu)成(B,j) -> (A,i)這么一條邊來,如此結(jié)點(diǎn)A的執(zhí)行就對(duì)結(jié)點(diǎn)B構(gòu)成依賴。

在Tensorflow的計(jì)算圖中,一般會(huì)包含兩個(gè)特殊的結(jié)點(diǎn)分別為Source節(jié)點(diǎn)(也稱Start節(jié)點(diǎn))與Sink節(jié)點(diǎn)(也稱為Finish節(jié)點(diǎn))。其中Source節(jié)點(diǎn)表示此節(jié)點(diǎn)不依賴于任何其它節(jié)點(diǎn)作為其輸入,而Sink節(jié)點(diǎn)則表示該節(jié)點(diǎn)并無任何輸出來作為其它節(jié)點(diǎn)的輸入。

Tensorflow計(jì)算圖示例

class graph代碼實(shí)例

  • Tensorflow中Graph構(gòu)造的描述可見于class Graph當(dāng)中(可在core/graph/graph.h中找到其定義)
class Graph {
 public:
  // Constructs a graph with a single SOURCE (always id kSourceId) and a
  // single SINK (always id kSinkId) node, and an edge from SOURCE->SINK.
  //
  // The graph can hold ops found in registry. `registry`s lifetime must be at
  // least that of the constructed graph's.
  explicit Graph(const OpRegistryInterface* registry);

  // Constructs a graph with a single SOURCE (always id kSourceId) and a
  // single SINK (always id kSinkId) node, and an edge from SOURCE->SINK.
  //
  // The graph can hold ops found in `flib_def`. Unlike the constructor taking
  // an OpRegistryInterface, this constructor copies the function definitions in
  // `flib_def` so its lifetime may be shorter than that of the graph's. The
  // OpRegistryInterface backing `flib_def` must still have the lifetime of the
  // graph though.
  explicit Graph(const FunctionLibraryDefinition& flib_def);

以上構(gòu)造函數(shù)當(dāng)中,我們看到Graph只需引入一個(gè)參數(shù)即OpRegistryInterface或FunctionLibraryDefinition。這兩個(gè)參數(shù)提供了具體每個(gè)節(jié)點(diǎn)的實(shí)際執(zhí)行定義。在我們構(gòu)建計(jì)算圖的時(shí)候,我們找到一個(gè)node的nodeDef(通常是基于google protocol buffer協(xié)議的node參數(shù)定義)后,會(huì)在OpRegistryInterface或FunctionLibraryDefinition當(dāng)中去獲取其具體的類型實(shí)現(xiàn)。也就是說我們?nèi)绻麑?shí)現(xiàn)了一個(gè)在某硬件平臺(tái)上優(yōu)化過了的Op或一種嶄新的Op操作,為了將此操作能夠作為計(jì)算圖的一個(gè)節(jié)點(diǎn)為我們的模型所用,那么需要將此新創(chuàng)建的Op實(shí)現(xiàn)函數(shù)注冊(cè)于OpRegistryInterface或FunctionLibraryDefinition結(jié)構(gòu)當(dāng)中。

  • 計(jì)算圖中有對(duì)Node與Edge的Add/Remove/Update等操作

其函數(shù)接口如下。具體的定義可見于core/graph/graph.cc當(dāng)中。本身實(shí)現(xiàn)起來因?yàn)槭鞘褂昧酥羔樻湵斫Y(jié)構(gòu)的DAG所以還是比較簡(jiǎn)單、容易理解的,在此就不多說了。

 // Adds a new node to this graph, and returns it. Infers the Op and
  // input/output types for the node. *this owns the returned instance.
  // Returns nullptr and sets *status on error.
  Node* AddNode(const NodeDef& node_def, Status* status);

  // Copies *node, which may belong to another graph, to a new node,
  // which is returned.  Does not copy any edges.  *this owns the
  // returned instance.
  Node* CopyNode(const Node* node);

  // Removes a node from this graph, including all edges from or to it.
  // *node should not be accessed after calling this function.
  // REQUIRES: node->IsOp()
  void RemoveNode(Node* node);

  // Adds an edge that connects the xth output of `source` to the yth input of
  // `dest` and returns it. Does not update dest's NodeDef.
  const Edge* AddEdge(Node* source, int x, Node* dest, int y);
  
  // Removes edge from the graph. Does not update the destination node's
  // NodeDef.
  // REQUIRES: The edge must exist.
  void RemoveEdge(const Edge* edge);
  // Updates the input to a node.  The existing edge to `dst` is removed and an
  // edge from `new_src` to `dst` is created. The NodeDef associated with `dst`
  // is also updated.
  Status UpdateEdge(Node* new_src, int new_src_index, Node* dst, int dst_index);
  • 圖之上的Op函數(shù)庫

class graph中有一個(gè)類成員為 FunctionLibraryDefinition ops_,其中包含了所有已知的具體類型的Op函數(shù)定義。而我們可利用以下函數(shù)來增加、拓展其Op函數(shù)庫。

// Adds the function and gradient definitions in `fdef_lib` to this graph's op
  // registry. Ignores duplicate functions, and returns a bad status if an
  // imported function differs from an existing function or op with the same
  // name.
  Status AddFunctionLibrary(const FunctionDefLibrary& fdef_lib);
  • Node節(jié)點(diǎn)對(duì)應(yīng)的宿主設(shè)備

Tensorflow當(dāng)中計(jì)算圖的執(zhí)行是并發(fā)的。圖上的每個(gè)Node都可被分布在不同的計(jì)算設(shè)備上計(jì)算。TF有提供API可以讓我們指定某個(gè)Op操作的宿主設(shè)備。當(dāng)然也有函數(shù)用來提供相應(yīng)的查詢操作。如下所示,見名可知其義。

const string& get_assigned_device_name(const Node& node) const {
    return device_names_[node.assigned_device_name_index()];
  }

  void set_assigned_device_name_index(Node* node, int device_name_index) {
    CheckDeviceNameIndex(device_name_index);
    node->assigned_device_name_index_ = device_name_index;
  }

  void set_assigned_device_name(Node* node, const string& device_name) {
    node->assigned_device_name_index_ = InternDeviceName(device_name);
  }

參考文獻(xiàn)

最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
平臺(tái)聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡(jiǎn)書系信息發(fā)布平臺(tái),僅提供信息存儲(chǔ)服務(wù)。
  • 序言:七十年代末,一起剝皮案震驚了整個(gè)濱河市,隨后出現(xiàn)的幾起案子,更是在濱河造成了極大的恐慌,老刑警劉巖,帶你破解...
    沈念sama閱讀 229,117評(píng)論 6 537
  • 序言:濱河連續(xù)發(fā)生了三起死亡事件,死亡現(xiàn)場(chǎng)離奇詭異,居然都是意外死亡,警方通過查閱死者的電腦和手機(jī),發(fā)現(xiàn)死者居然都...
    沈念sama閱讀 98,860評(píng)論 3 423
  • 文/潘曉璐 我一進(jìn)店門,熙熙樓的掌柜王于貴愁眉苦臉地迎上來,“玉大人,你說我怎么就攤上這事。” “怎么了?”我有些...
    開封第一講書人閱讀 177,128評(píng)論 0 381
  • 文/不壞的土叔 我叫張陵,是天一觀的道長(zhǎng)。 經(jīng)常有香客問我,道長(zhǎng),這世上最難降的妖魔是什么? 我笑而不...
    開封第一講書人閱讀 63,291評(píng)論 1 315
  • 正文 為了忘掉前任,我火速辦了婚禮,結(jié)果婚禮上,老公的妹妹穿的比我還像新娘。我一直安慰自己,他們只是感情好,可當(dāng)我...
    茶點(diǎn)故事閱讀 72,025評(píng)論 6 410
  • 文/花漫 我一把揭開白布。 她就那樣靜靜地躺著,像睡著了一般。 火紅的嫁衣襯著肌膚如雪。 梳的紋絲不亂的頭發(fā)上,一...
    開封第一講書人閱讀 55,421評(píng)論 1 324
  • 那天,我揣著相機(jī)與錄音,去河邊找鬼。 笑死,一個(gè)胖子當(dāng)著我的面吹牛,可吹牛的內(nèi)容都是我干的。 我是一名探鬼主播,決...
    沈念sama閱讀 43,477評(píng)論 3 444
  • 文/蒼蘭香墨 我猛地睜開眼,長(zhǎng)吁一口氣:“原來是場(chǎng)噩夢(mèng)啊……” “哼!你這毒婦竟也來了?” 一聲冷哼從身側(cè)響起,我...
    開封第一講書人閱讀 42,642評(píng)論 0 289
  • 序言:老撾萬榮一對(duì)情侶失蹤,失蹤者是張志新(化名)和其女友劉穎,沒想到半個(gè)月后,有當(dāng)?shù)厝嗽跇淞掷锇l(fā)現(xiàn)了一具尸體,經(jīng)...
    沈念sama閱讀 49,177評(píng)論 1 335
  • 正文 獨(dú)居荒郊野嶺守林人離奇死亡,尸身上長(zhǎng)有42處帶血的膿包…… 初始之章·張勛 以下內(nèi)容為張勛視角 年9月15日...
    茶點(diǎn)故事閱讀 40,970評(píng)論 3 356
  • 正文 我和宋清朗相戀三年,在試婚紗的時(shí)候發(fā)現(xiàn)自己被綠了。 大學(xué)時(shí)的朋友給我發(fā)了我未婚夫和他白月光在一起吃飯的照片。...
    茶點(diǎn)故事閱讀 43,157評(píng)論 1 371
  • 序言:一個(gè)原本活蹦亂跳的男人離奇死亡,死狀恐怖,靈堂內(nèi)的尸體忽然破棺而出,到底是詐尸還是另有隱情,我是刑警寧澤,帶...
    沈念sama閱讀 38,717評(píng)論 5 362
  • 正文 年R本政府宣布,位于F島的核電站,受9級(jí)特大地震影響,放射性物質(zhì)發(fā)生泄漏。R本人自食惡果不足惜,卻給世界環(huán)境...
    茶點(diǎn)故事閱讀 44,410評(píng)論 3 347
  • 文/蒙蒙 一、第九天 我趴在偏房一處隱蔽的房頂上張望。 院中可真熱鬧,春花似錦、人聲如沸。這莊子的主人今日做“春日...
    開封第一講書人閱讀 34,821評(píng)論 0 28
  • 文/蒼蘭香墨 我抬頭看了看天上的太陽。三九已至,卻和暖如春,著一層夾襖步出監(jiān)牢的瞬間,已是汗流浹背。 一陣腳步聲響...
    開封第一講書人閱讀 36,053評(píng)論 1 289
  • 我被黑心中介騙來泰國(guó)打工, 沒想到剛下飛機(jī)就差點(diǎn)兒被人妖公主榨干…… 1. 我叫王不留,地道東北人。 一個(gè)月前我還...
    沈念sama閱讀 51,896評(píng)論 3 395
  • 正文 我出身青樓,卻偏偏與公主長(zhǎng)得像,于是被迫代替她去往敵國(guó)和親。 傳聞我的和親對(duì)象是個(gè)殘疾皇子,可洞房花燭夜當(dāng)晚...
    茶點(diǎn)故事閱讀 48,157評(píng)論 2 375

推薦閱讀更多精彩內(nèi)容