本文作者:李杰
TF计算图从逻辑层来讲,由op与tensor构成。op是项点代表计算单元,tensor是边代表op之间流动的数据内容,两者配合以数据流图的形式来表达计算图。那么op对应的物理层实现是什么?TF中有哪些op,以及各自的适用场景是什么?op到底是如何运行的?接下来让我们一起探索和回答这些问题。
一、初识op
1.1op定义
op代表计算图中的节点,是tf.Operation对象,代表一个计算单元。用户在创建模型和训练代码时,会创建一系列op及其依赖关系,并将这些op和依赖添加到tf.Graph对象中(一般为默认图)。比如:tf.matmul()就是一个op,它有两个输入tensor和一个输出tensor。
1.2op分类
op的分类一般有多个视角,比如按是否内置划分、按工作类型划分。
按是否内置划分,一般分为:内置op和自定义op(见“二、自定义op”部分介绍)。
按工作类型划分,一般分为:常见数学op、数组op、矩阵op、有状态op、神经网络op、检查点op、队列与同步op、控制流op。TF白皮书对内置op的分类总结如下:
1.3op与kernel
op一般都有名称且代表一个抽象的计算过程。op可以设置若干属性,但这些属性必须在编译期提供或推理得到,因为它们用来实例化一个节点对象从而执行真正的计算。属性的经典用法就是拿来支持类型多态,比如两个浮点张量的矩阵乘法与两个整型张量的矩阵乘法。
kernel是op在指定设备类型(CPU/GPU)上的具体实现。TF二进制库通过注册机制定义了一系列op及对应的kernel实现,用户可以提供额外的op定义与kernel实现进行扩充。一般来说,一个op对应多个kernel实现。
接下来让我们一起用矩阵乘法MatMul算子的相关代码来理解op与kernel的关系(此处不必纠结代码细节,只需体会op与kernel关系即可):
//首先给出op注册的定义。其中输入输出支持泛型,其合法类型在Attr中进行枚举。//代码位置tensorflow1.15.5\tensorflow\core\ops\math_ops.ccREGISTER_OP("MatMul").Input("a:T").Input("b:T").Output("product:T").Attr("transpose_a:bool=false").Attr("transpose_b:bool=false").Attr("T:{bfloat16,half,float,double,int32,int64,