TensorFlow架構與設計:圖模塊

計算圖是TensorFlow領域模型的核心。本文通過對計算圖領域模型的梳理,講述計算圖構造的基本原理。

Edge持有前驅節點與后驅節點,從而實現了計算圖的連接,也是計算圖前向遍歷,后向遍歷的銜接點。

邊上的數據以Tensor的形式傳遞,Tensor的標識由源節點的名稱,及其所在邊的src_output唯一確定。也就是說,tensor_id = op_name:src_output

src_output與dst_input

Edge持有兩個重要的屬性:

  • src_output:表示該邊為前驅節點的第src_output條輸出邊;
  • dst_input:表示該邊為后驅節點的第dst_input條輸入邊。

例如,存在兩個前驅節點s1, s2,都存在兩條輸出邊;存在兩個后驅節點d1, d2,都存在兩條輸入邊。

邊索引

控制依賴

計算圖中存在兩類邊,

  • 普通邊:用于承載Tensor,常用實線表示;
  • 控制依賴:控制節點的執行順序,常用虛線表示。

特殊地,控制依賴邊,其src_output, dst_input都為-1(Graph::kControlSlot),暗喻控制依賴邊不承載任何數據,僅僅表示計算的依賴關系。

bool Edge::IsControlEdge() const {
   return src_output_ == Graph::kControlSlot;
}

節點

Node(節點)持有零條或多條輸入/輸出的邊,分別使用in_edges, out_edges表示。另外,Node持有NodeDef, OpDef。其中,NodeDef持有設備分配信息,及其OP的屬性值集合;OpDef持有OP的元數據。

節點

輸入邊

在輸入邊的集合中按照索引線性查找,當節點輸入的邊比較多時,可能會成為性能的瓶頸。依次類推,按照索引查找輸出邊,算法相同。

Status Node::input_edge(int idx, const Edge** e) const {
  for (auto edge : in_edges()) {
    if (edge->dst_input() == idx) {
      *e = edge;
      return Status::OK();
    }
  }
  return errors::NotFound("not found input edge ", idx);
}

前驅節點

首先通過idx索引找到輸入邊,然后通過輸入邊找到前驅節點。依次類推,按照索引查找后驅節點,算法相同。

Status Node::input_node(int idx, const Node** n) const {
  const Edge* e;
  TF_RETURN_IF_ERROR(input_edge(idx, &e));
  if (e == nullptr) {
    *n = nullptr;
  } else {
    *n = e->src();
  }
  return Status::OK();
}

Graph(計算圖)就是節點與邊的集合,領域模型何其簡單。計算圖是一個DAG圖,計算圖的執行過程將按照DAG的拓撲排序,依次啟動OP的運算。其中,如果存在多個入度為0的節點,TensorFlow運行時可以實現并發,同時執行多個OP的運算,提高執行效率。

空圖

計算圖的初始狀態,并非是一個空圖。實現添加了兩個特殊的節點:Source與Sink節點,分別表示DAG圖的起始節點與終止節點。其中,Source的id為0,Sink的id為1;依次論斷,普通OP節點的id將大于1。

另外,Source與Sink之間,通過連接「控制依賴」的邊,保證計算圖的執行始于Source節點,終于Sink節點。它們之前連接的控制依賴邊,其src_output, dst_input值都為-1。

習慣上,僅包含Source與Sink節點的計算圖也常常稱為空圖。

空圖
Node* Graph::AddEndpoint(const char* name, int id) {
  NodeDef def;
  def.set_name(name);
  def.set_op("NoOp");

  Status status;
  Node* node = AddNode(def, &status);
  TF_CHECK_OK(status);
  CHECK_EQ(node->id(), node_id);
  return node;
}

Graph::Graph(const OpRegistryInterface* ops)
    : ops_(ops), arena_(8 << 10 /* 8kB */) {
  auto src  = AddEndpoint("_SOURCE", kSourceId);
  auto sink = AddEndpoint("_SINK",   kSinkId);
  AddControlEdge(src, sink);
}

非空圖

在前端,用戶使用OP構造器,將構造任意復雜度的計算圖。對于運行時,無非就是將用戶構造的計算圖通過控制依賴的邊與Source/Sink節點連接,保證計算圖執行始于Source節點,終于Sink節點。

非空圖

添加邊

計算圖的構造過程非常簡單,首先通過Graph::AddNode在圖中放置節點,然后再通過Graph::AddEdge在圖中放置邊,實現節點之間的連接。

const Edge* Graph::AllocEdge() const {
  Edge* e = nullptr;
  if (free_edges_.empty()) {
    e = new (arena_.Alloc(sizeof(Edge))) Edge;
  } else {
    e = free_edges_.back();
    free_edges_.pop_back();
  }
  e->id_ = edges_.size();
  return e;
}

const Edge* Graph::AddEdge(Node* source, int x, Node* dest, int y) {
  auto e = AllocEdge();
  e->src_ = source;
  e->dst_ = dest;
  e->src_output_ = x;
  e->dst_input_ = y;

  CHECK(source->out_edges_.insert(e).second);
  CHECK(dest->in_edges_.insert(e).second);

  edges_.push_back(e);
  edge_set_.insert(e);
  return e;
}

添加控制依賴邊,則可以轉發調用Graph::AddEdge實現。

const Edge* Graph::AddControlEdge(Node* src, Node* dst) {
  return AddEdge(src, kControlSlot, dst, kControlSlot);
}

開源技術書

https://github.com/horance-liu/tensorflow-internals
最后編輯于
?著作權歸作者所有,轉載或內容合作請聯系作者
平臺聲明:文章內容(如有圖片或視頻亦包括在內)由作者上傳并發布,文章內容僅代表作者本人觀點,簡書系信息發布平臺,僅提供信息存儲服務。

推薦閱讀更多精彩內容