tensorflow源碼解析之framework-node

目錄

1. 什么是node
2. node_def
3. 關系圖
4. 涉及的文件
5. 迭代記錄

1. 什么是node

TF中的計算圖由節點組成，每個節點包含了一個操作，表示這個節點的作用，比如，如果一個節點的作用是做矩陣乘法，那么它的輸入是兩個矩陣，輸出是兩個輸入矩陣相乘的結果。節點是自帶結構的，每個節點都包含了輸入的來源，因此若干節點的集合就能無需其它信息的生成一張圖。節點必須被放置在某一個設備上，為了減少跨設備傳輸數據，也為了提高計算效率，節點的放置是一個受限條件下的優化問題，為此TF還專門開發了相應的節點放置算法。

2. node_def

我們先來看一下節點的定義：

message NodeDef {
    string name = 1;//節點名稱
    string op = 2;//節點包含的操作名稱
    repeated string input = 3;//節點的輸入，詳見正文
    string device = 4;//節點所在的設備
    map<string, AttrValue> attr = 5;//節點包含的op，其參數的具體賦值
};

由於節點對於后續理解graph和kernel非常重要，因此我們詳細分析一下它包含的字段：

• 關於input，剛才提到，節點是自帶結構信息的，這個信息就包含在input字段中。input的格式是node:src_output，其中的node表示輸入的節點名稱，src_output表示該輸入是輸入節點的第幾個輸出。這樣聽起來可能很繞，舉個例子，假設node1和node2是兩個節點，它們分別產生兩個輸出，node1:0，node1:1，node2:0，node2:1。現在我們有第三個節點node3，它需要來自前兩個節點的輸出作為輸入，具體需要的輸入是node1:1和node2:0，那么，node3的input字段就是node1:1和node2:0，如下圖所示。另外，如果需要的是輸入節點的第1個輸出，那么后面的端口號可以忽略，比如node2:0可以直接寫為node2。

graph LR node1-->|輸出1|node1:0 node1-->|輸出2|node1:1 node2-->|輸出1|node2:0 node2-->|輸出2|node2:1 node1:1-->node3 node2:0-->node3

• 關於input，再啰嗦一點，節點的輸入可以是真實輸入，也可以是控制輸入。控制輸入表示節點並不需要這個輸入的具體數值做計算，僅僅是為了說明一種計算的先后關系，即當前節點的計算，必須在所有的控制輸入計算都完成之后進行。在表示input字段時，控制輸入需要跟在常規輸入的后面，並且用^node表示。
• 關於device，它表示了我們對於當前節點可以被放置的設備的需求。這個描述可以是完全限定的，比如"/job:worker/replica:0/task:1/gpu:3"，也可以是局部限定的，比如"/job:worker/gpu:3"，也可以不限定。因此這個字段指代的，可能不是某一個具體的設備，這就需要上文提到的節點放置算法，對節點所在的設備進行選擇了。也就是說，device這個字段有些情況下只是給節點所在的設備限制了一個粗略的條件，節點放置算法需要在考慮這些條件的基礎上，為節點選擇合適的設備。關於節點放置算法的細節，我們將在后面詳述。
• 關於attr，我們前面講過，op是一個抽象信息，這個抽象信息包含了很多未賦值的參數，在運行時構建圖的時候，這些參數就一定要被賦值了，否則圖計算就無法進行。因此這里必須包含op中所有未賦值的參數的具體數值。如果op中的參數有默認值，這里可以對其進行覆蓋，也可以忽略，這樣該參數就會使用默認值。

按照慣例，TF會給每一個核心概念設計一個構建類，node也不例外，這個類就是NodeDefBuilder。我們先來看一下它的用法：

NodeDef node_def;
Status status = NodeDefBuilder(node_name, op_name)
                    .Input(...)
                    .Attr(...)
                    .Finalize(&node_def);

可見，與OpDefBuilder類似，我們也可以采用鏈式規則，通過逐個設置NodeDef的字段來構建它。下面我們看下NodeDefBuilder包含的私有數據：

class NodeDefBuilder {
  private:
    const OpDef* op_def_;
    NodeDef node_def_;
    int inputs_specified_;
    std::vector<string> control_inputs_;
    std::vector<string> errors_;
};

這里之所以會包含OpDef，是因為在NodeDef中，僅包含了操作的名稱。

3. 關系圖

graph LR NodeDefBuilder-->|構建|NodeDef NodeDef-.包含.->節點名稱name NodeDef-.包含.->節點包含的操作名稱op NodeDef-.包含.->輸入input NodeDef-.包含.->節點所在設備名稱device NodeDef-.包含.->節點屬性也是操作屬性attr

4. 涉及的文件

• node_def

5. 迭代記錄

• v1.0 2018-08-27 文檔創建
• v2.0 2018-09-09 文檔重構

github地址