Tensorflow核心代码解析之计算图篇其一:计算图结构初探

介绍

当今计算机科学给人对未来最大想象的莫过于人工智能的大规模应用前景。它对于人类文明进步所带来的潜在贡献可以被视为与四大发明、蒸汽机、电力、计算机等人类史标杆性工具具有同等的地位。当今人工智能的蓬勃发展主要是机器学习尤其是深度学习的大规模成功应用。凭着日益增多的海量数据,快速发展的计算机并行计算能力,快速迭代、高效更新的各种模型、算法、策略以及各个国家、政府、大企业对它的日益重视与超高的资金、政策投入,当今AI的发展速度真可谓一日千里!

深度学习最终要与具体行业场景有效率地结合起来才能发挥出其效益来。当下的整个深度学习已经行成了良好的产业链体系。最下端的是用于深度学习加速的各种硬件芯片如CPU/GPU/FPGA/ASIC专用芯片等。目前此领域里面GPU凭其优良的超多弱核心并行计算能力独领风骚;但CPU在推理加速、成本等方面也挺有竞争力;FPGA凭其灵活性也非常适宜于进行AI方案的原型设计,但因其开发难度较大,生态相对较缺乏,当下大公司里面大规模部署应用FPGA的唯有微软;至于ASIC专用芯片,可谓是给了诸多有心在AI半导体上面实现弯道超车的公司一个很好的机会,尤其是那些有着大规模机器学习用户,可以基于上层封装提供AI云计算实例服务的公司如Google就声势浩大地推出了自己快速迭代着的TPU,号称常规模型(如Resnet-50)加速快于最新的GPU,同时功耗更为节省,其它像Amazon,Facebook,Ali等云计算公司也在搞自己的AI专用ASIC芯片。此外一些创业型公司也对此雄心勃勃,国内已经涌现出了一堆此类的独角兽像卖自家矿机发了大财的Bit大陆,技术势力雄厚的地平线等等。还有些手机大厂则将精力用在了终端一侧的AI芯片研发像华为、苹果等已经有了自己的AI芯片并部署在了自己的手机新品当中,其它像小米、三星等也在纷纷跟进。需要提下的是ARM这个在移动时代的主要得意者也于最近发布了自己家设计的用于AI加速的各种硬件IP。

AI芯片

计算芯片向上走则是将一些基本运算如矩阵乘积、各类型卷积运算等结合硬件平台优化过了的数学计算库如用于Intel CPU端的MKL/MKLDNN,用于nVidia GPU的CUDA/cuDNN,用于FPGA专用DNN网络加速过的openCL,还有针对各大ASIC芯片产商针对自己ASIC加速过了的种种计算库套件。这些数学计算库基本由AI芯片产商自己来完成,目的即在于借用软件的力量给自家的硬件以强大的驱动力。一般他们会选择将优化过了的核心程序开源、公布出来为自己的客户所借鉴、使用,最终通过AI芯片的出售来获得价值。值得一提的是nVIDIA在基于CUDA/cuDNN上面的多年耕耘,相关社区、生态的耐心培养直接带来了它们今天的巨大成功。现在半导体公司已经再不同往日只需要卖出芯片即可了。芯片相关的软件库(并行计算库等)的性能,对用户提供API的友好性,用户社区的培养,用户支持的力度等等软实力真正是愈益重要。

底层计算库再向上走则是使用这些优化过了的数学计算库来完成基本类型计算,然后将之抽象封装后向上提供出友好用户API的深度学习框架。这些深度学习框架是我们软件人员开发应用的基本工具。当前最流行的框架有Google的Tensorflow,Facebook的Pytorch,Amazon的Mxnet,微软的CNTK,当然还有传统社区在维护的bvlc/Caffe等等。它们大都提供类似的功能,相似的API用于用户程序构建计算图,并能将图方便地导入、导出为序列化文件,还提供了基于Framework level对图的一些优化如合并(fusion),去重(典型的如CSE),并行计算(通过使用OMP等并行库)等,此外还有一些功能如用于进行内存分配、管理及线程执行、调度、检测的session/workspace等,当然还有用于具体执行某计算的op / kernel等,这也是常规计算优化的核心所在。

在这些计算框架中,无疑Google brain团队开发的Tensorflow是最为流行的。它的框架设计最为复杂,可以天生地支持模型并行训练、推理等,它的背后有一个google强大的开发团队在快速迭代、开发,它的底下也有集成当前最好的像cuDNN/mklDNN/TensorRT等加速技术,它的用户社区也已经非常完善、活跃(作为程序员这个还是蛮重要的。毕竟在APP开发中出了问题,肯定都希望通过在网上翻一下看有没有人踩过类似的坑以来快速解决问题。)。

当下对于如何使用Tensorflow来开发一个AI程序,构建深度学习模型并进行训练或推理的文章已经很多了。本系列单元中笔者想试着跟大家一起理一下它框架核心的一些代码实现。无益它对于我们基于TF做一些开发,加深对TF的理解是很有帮助的。此外TF框架的设计、代码实现非常良好,对它们的理解、梳理清楚对于我们日常的软件设计、开发也会有较强的工程借鉴意义。

计算图

计算图(Graph)描述了一组需要依次序完成的计算单元以及表示这些计算单元之间相互依赖的关系。一般的深度学习模型都会被分化组装成一个单向无环图(DAG)来执行。图当中的结点(node)用来表示某一具体的计算单元(如Multmul结点表示两个张量之间的乘积,Conv结点则表示两个张量之间的卷积计算)。图上的片(edge)则被用来表示两个结点之间的依赖关系。比如A结点的第i个输出来自于B结点的第j输入,那么就会构成(B,j) -> (A,i)这么一条边来,如此结点A的执行就对结点B构成依赖。

在Tensorflow的计算图中,一般会包含两个特殊的结点分别为Source节点(也称Start节点)与Sink节点(也称为Finish节点)。其中Source节点表示此节点不依赖于任何其它节点作为其输入,而Sink节点则表示该节点并无任何输出来作为其它节点的输入。

Tensorflow计算图示例

class graph代码实例

  • Tensorflow中Graph构造的描述可见于class Graph当中(可在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);

以上构造函数当中,我们看到Graph只需引入一个参数即OpRegistryInterface或FunctionLibraryDefinition。这两个参数提供了具体每个节点的实际执行定义。在我们构建计算图的时候,我们找到一个node的nodeDef(通常是基于google protocol buffer协议的node参数定义)后,会在OpRegistryInterface或FunctionLibraryDefinition当中去获取其具体的类型实现。也就是说我们如果实现了一个在某硬件平台上优化过了的Op或一种崭新的Op操作,为了将此操作能够作为计算图的一个节点为我们的模型所用,那么需要将此新创建的Op实现函数注册于OpRegistryInterface或FunctionLibraryDefinition结构当中。

  • 计算图中有对Node与Edge的Add/Remove/Update等操作

其函数接口如下。具体的定义可见于core/graph/graph.cc当中。本身实现起来因为是使用了指针链表结构的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函数库

class graph中有一个类成员为 FunctionLibraryDefinition ops_,其中包含了所有已知的具体类型的Op函数定义。而我们可利用以下函数来增加、拓展其Op函数库。

// 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节点对应的宿主设备

Tensorflow当中计算图的执行是并发的。图上的每个Node都可被分布在不同的计算设备上计算。TF有提供API可以让我们指定某个Op操作的宿主设备。当然也有函数用来提供相应的查询操作。如下所示,见名可知其义。

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);
  }

参考文献

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 151,511评论 1 330
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 64,495评论 1 273
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 101,595评论 0 225
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 42,558评论 0 190
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 50,715评论 3 270
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 39,672评论 1 192
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,112评论 2 291
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 29,837评论 0 181
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 33,417评论 0 228
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 29,928评论 2 232
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 31,316评论 1 242
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 27,773评论 2 234
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 32,253评论 3 220
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 25,827评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,440评论 0 180
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 34,523评论 2 249
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 34,583评论 2 249

推荐阅读更多精彩内容