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

介紹

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

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

AI芯片

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

底層計算庫再向上走則是使用這些優(yōu)化過了的數(shù)學(xué)計算庫來完成基本類型計算，然后將之抽象封裝后向上提供出友好用戶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)建計算圖，并能將圖方便地導(dǎo)入、導(dǎo)出為序列化文件，還提供了基于Framework level對圖的一些優(yōu)化如合并(fusion)，去重（典型的如CSE），并行計算（通過使用OMP等并行庫）等，此外還有一些功能如用于進(jìn)行內(nèi)存分配、管理及線程執(zhí)行、調(diào)度、檢測的session/workspace等，當(dāng)然還有用于具體執(zhí)行某計算的op / kernel等，這也是常規(guī)計算優(yōu)化的核心所在。

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

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

計算圖

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

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

Tensorflow計算圖示例

class graph代碼實例

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

  • 計算圖中有對Node與Edge的Add/Remove/Update等操作

其函數(shù)接口如下。具體的定義可見于core/graph/graph.cc當(dāng)中。本身實現(xiàn)起來因為是使用了指針鏈表結(jié)構(gòu)的DAG所以還是比較簡單、容易理解的，在此就不多說了。

 // 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中有一個類成員為 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é)點對應(yīng)的宿主設(shè)備

Tensorflow當(dāng)中計算圖的執(zhí)行是并發(fā)的。圖上的每個Node都可被分布在不同的計算設(shè)備上計算。TF有提供API可以讓我們指定某個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)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點，簡書系信息發(fā)布平臺，僅提供信息存儲服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容