最近在读《TensorFlow 内核剖析》这本书,作者刘光聪。本篇文章摘自 TensorFlow架构与设计:OP本质论 - 简书

符号编程

TensorFlow的计算过程是一个延迟计算,是一种典型的基于符号的编程范式。从计算时间轴看,计算过程基本分为2个阶段:

  • 图构造期:负责计算图的构造;
  • 图执行期:负责计算图的执行。

其中,在系统初始化时,系统实现对所有OP进行扫描注册,并保存于OpRegistry之中。

注册OP

理论上,OP的注册发生在系统初始化阶段。后端系统,可以使用REGISTER_OP实用宏注册OP。前端系统,也存在类似的OP注册机制。

使用REGISTER_OP注册OP过程,实际上是一个REGISTER_OP描述到OpDef表示的翻译过程。OpDefBuilder通过链式调用Input, Output, Attr方法分别构造OP的输入、输出列表,及其属性列表。最后,通过调用Finalize成员函数,经过解析字符串表示,将其翻译为OpDef的内在表示,最后注册到OpRegistry之中。

例如,REGISTER_OP("ZerosLike")向系统注册了一个zeros_like的OP,在运行时实现了OpDef的翻译表达。

构造OP

在前端,用户使用OP构造器实现OP的构造,并将OP注册到计算图中。在计算图构造期间,OP的输入/输出的类型,Shape得以确定,OP属性值也得以确定。

计算图的构造过程,实际上就是GraphDef定义过程。其中,OP的属性值承载于NodeDef,计算图构造期间,NodeDef的属性值得以确定。

在计算图执行启动时,通过调用Session.run,将整个GraphDef传递给后端,并启动计算图的执行。

执行OP

在计算图执行期间,输入由上游OP流入得以确定,根据特定设备类型,输入输出类型,多态选择合适的Kernel实现,并启动Kernel的计算过程。

参考资料

TensorFlow架构与设计:OP本质论 - 简书