tensorflow核心概念之op

Posted 2020-12-31 jicanghai

    TF中的op代表了对“操作”的抽象，说它抽象是因为，op仅仅说明了操作是做什么用的，但没有说明具体怎么做。举个例子，MatMul是一个操作，它表示了矩阵乘法，但并不包含矩阵乘法的具体实现，因为我们知道，在CPU和GPU上，矩阵乘法的高效率实现是完全不同的。为了把不同设备上的实现细节隐藏起来，为相同的计算提供统一的对外表示，TF提出了op的概念。

    为了方便序列化，TF中很多核心概念定义在proto文件中，操作的定义OpDef放在op.proto（/tensorflow/core/framework/op.proto）文件中。因此，OpDef仅包含一些静态的数值信息，比如操作的名称，输入、输出类型，参数，描述，以及一些计算相关的属性信息（是否可交换、是否可聚集、是否有状态）。随着对TF剖析的深入，我们会发现核心概念往往都有xx_def定义，比如graph_def，node_def，kernel_def等。

    仅有静态信息是不够的，我们在构建一个op的时候，还需要灵活的对op_def中的属性进行设置，于是就有了对OpDefBuilder类的需求。后续我们还会看到针对其它核心概念的构建类，比如graph_builder，node_builder等等。

    有了OpDefBuilder类之后，我们构建一个op就比较简单了，首先创建一个OpDefBuilder的对象，然后依次调用该对象的属性set函数，直到完成构建。这个过程对于单个op来说相当轻松，但是当我们需要构建上百个op时，这个任务就会变得非常繁琐。一个想法是，如果能把OpDefBuilder对象属性设置的函数串联起来，形成一个chain，这个过程就会方便很多，于是就诞生了OpDefBuilderWrapper，它仅仅是OpDefBuilder的一个封装，使属性值设置的行为可以串联进行，简化了操作定义。

    TF的ops目录下存放了上百个op定义，但并不是每一个TF应用都需要所有的定义。为了为TF底层库进行减负，我们需要把不需要的op省略掉，为了对此作出支持，OpDefBuilderWrapper被设计成为模板类，我们可以通过宏定义选择哪些操作被编译进TF，哪些可以被省略。

    TF官方定义了大量的op，用户本身也会有大量的自定义op，为了对op进行集中管理，TF提供了OpRegistryInterface接口，它只包含一个函数，Lookup，通过操作名称查找OpDef。具体由OpRegistry和OpListRegistry两个类来实现，我们主要关注前者。OpRegistry内部包含了一个op名称到op_def的映射，在注册的时候，采取了延迟注册的机制。对于申请注册的OpDef，先放入一个deferred_的向量，在下次需要对已注册OpDef进行查找之前，将deferred_中的操作注册，并清空deferred_，这样既能保证所有的操作都被注册，也能保证注册过程的高效。