探索TensorFlow核心组件系列之Operation的底层机制

TensorFlow中，Operation（操作）是构成计算图的核心组件之一。它代表了计算图中的一个节点，执行着各种数学运算、数据处理和变换操作。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

其中，每个Operation都具有自己的计算逻辑和属性，用于描述它所执行的具体操作。操作可以是简单的数学运算，例如加法和乘法，也可以是复杂的神经网络层或自定义的算法。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

Operation的底层机制包括以下关键要素：计算函数、输入输出张量、属性和内存管理。这些要素共同协作，使得Operation能够有效地执行计算任务并产生结果。可以说，Operation的底层机制是TensorFlow实现高效计算和灵活模型构建的关键所在。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

TensorFlow中Op就类似于Spark中的计算节点（或称为Node），负责某种抽象计算，它表示计算图中的节点。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

但相比Spark中Node计算节点的分类，TF中就复杂的多，主要包括：文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

数学运算 OP，主要进行加、减、乘、除、矩阵乘法、矩阵转置等运算；文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

神经网络 OP，主要进行神经网络模型，包括卷积、池化、全连接等；文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

数据操作 OP，主要对输入和输出数据进行处理；文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

优化器 OP，主要进行模型训练过程中的优化操作，如梯度下降、自适应学习率优化算法；文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

控制流 OP，主要用于控制程序运行流程，如条件语句、循环语句等；文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

还有图像操作，分布式，数据集，特定领域的Op。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

我们都知道在TensorFlow中，前端负责构图，后端负责运行，那么分别来看下前后端的Op实现逻辑。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

1. Op前端（Python）实现

1.1 构建图从创建Op开始文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

在计算图的构建期间，是通过Op的构造器，创建Operation实例的，并在创建的过程中将其注册到默认图中。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

其中，Operation 类的元数据由 OpDef 与 NodeDef 持有，它们以 ProtoBuf 的格式存在，它描述 了 Operation 最重要的东西，也是最本质的东西。其中，OpDef 描述了 OP 的静态属性信息，例如 OP 的名称， 输入/输出参数列表，属性集定义等信息。而 NodeDef 描述了 OP 的动态属性值信息，例如 属性值等信息，后面我们再详细描述下它们。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

1.2 通过Operation构造器创建Op文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

在Python前端中，Op对象是由Operation构造器创建的，具体代码如下所示：文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

class Operation(object):
       def __init__(self, node_def, g, inputs=None, output_types=None,
                    control_inputs=None, input_types=None, original_op=None,
                    op_def=None):
				 # 1. NodeDef
         self._node_def = copy.deepcopy(node_def)
         # 2. OpDef
         self._op_def = op_def
         # 3. Graph
         self._graph = g
         # 4. Input types
         if input_types is None:
           input_types = [i.dtype.base_dtype for i in self._inputs]
         self._input_types = input_types
         # 5. Output types
         if output_types is None:
           output_types = []
         self._output_types = output_types 
         # 6. Inputs
         if inputs is None:
           inputs = []
         self._inputs = list(inputs)
         # 7. Control Inputs.
         if control_inputs is None:
           control_inputs = []
         self._control_inputs = []
         for c in control_inputs:
				   c_op = self._get_op_from(c)
           self._control_inputs.append(c_op)
         # 8. Outputs
         self._outputs = [Tensor(self, i, output_type)
                          for i, output_type in enumerate(output_types)]
         # 9. Build producter-consumer relation.
         for a in self._inputs:
           a._add_consumer(self)
         # 10. Allocate unique id for opeartion in graph.
         self._id_value = self._graph._next_id()

从上面的代码可以看出，创建一个Op对象的主要参数有，通过传入一个原Op作为待构造Operation的输入，传入Input的Tensor作为上游输入，从这里可以看出Op的上游并非直接关联Op，而是上游Op输出的Tensor，通过这种依赖关系在session.run()运行时找到最小的依赖子图。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

从Operation类的成员变量可以看出，其主要包含了OpDef和NodeDef两个主要成员变量，同时记录了Graph信息，输入输出类型， 输入Tensor列表与控制输入列表。并建立了生产者消费者关系，通过遍历了所有的上游输入Tensor，并将当前Operation加入了其维护的消费者列表中。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

下面我们来看下NodeDef的组成：文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

node {
      name: "a"
      op: "VariableV2"
      attr {
        key: "_output_shapes"
        value {
          list {
            shape {
            }
          }
        }
      }
      attr {
        key: "container"
        value {
          s: ""
        }
      }
      attr {
        key: "dtype"
        value {
          type: DT_FLOAT
        }
      }
      attr {
        key: "shape"
        value {
          shape {
          }
        }
      }
      attr {
        key: "shared_name"
        value {
          s: ""
        }
      }
    }

从上面可以看出NodeDef包含了，当前Node的Op 名字, 属性与输入等信息。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

需要注意的是，有些Operation的创建是不需要Input的，比如Constant，这样的Op被称为源Op。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

1.3 创建后的Op对象被作为一个Node节点加入Graph中，最后创建Session, 并运行Op。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

从该 OP 为末端反向遍历图，根据依赖关系，寻找最小依赖的子图，寻找依赖的过程中就是寻找其input，直到没有input为止，并在默认的 Session 中执行该子图。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

注意，第一次执行该图时，会为其分配Executor, 但如果图没改变，再次支持时，则会直接返回上次分配的Executor。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

前端系统中通过Operation来创建Op从而构建整个Graph, 然后将其序列化成 Protocol Buffer 格式（即 .pb 文件）转递到后端系统。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

2. Op后端（C++）实现

在C++后端系统中，Python前端Operation对应的是后端的Node实现。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

其中Node(节点) 可以拥有零条或多条输入/输出的边，并使用 in_edges, out_edges 分别表 示输入边和输出边的集合。另外，后端的Node对象也持有 NodeDef, OpDef。其中，NodeDef 包含设备分配 信息，及其 OP 的属性值列表; OpDef 持有 OP 的元数据，包括 OP 输入输出类型等信息。Node与Edge共同组成了Graph。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

class Node {
 public:
  string DebugString() const;
  int id() const { return id_; }
  const NodeDef& def() const;
  const OpDef& op_def() const;
  DataType input_type(int32 i) const;
  DataType output_type(int32 o) const;
  const string& requested_device() const;
  const EdgeSet& in_edges() const { return in_edges_; }
  const EdgeSet& out_edges() const { return out_edges_; }
  ...
  // Node type helpers.
  bool IsSource() const { return id() == 0; }
  bool IsSink() const { return id() == 1; }
  // Anything other than the special Source & Sink nodes.
  bool IsOp() const { return id() > 1; }

  Status input_edge(int idx, const Edge** e) const;

  Status input_node(int idx, const Node** n) const;
  Status input_tensor(int idx, OutputTensor* t) const;
 private:
  friend class Graph;
  Node();
  ...
};

从上面可以看出，后端的Node对象与前端Operation对象是类似的，其不仅持有 NodeDef, OpDef，也包括了in_edges, out_edges等构图信息。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

在大家印象里Op指的是计算函数，例如tf.divide，而目前介绍了一通图中的Node对象。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

那么Node对象是如何和具体的Op发生关系呢？文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

其实其主要靠的就是Node中的OpDef对象，下面来详细说下Op的定义组成部分。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

在TensorFlow中所有的Op定义主要包括以下几个部分，下面我们以ZerosLikeOp简单举例：文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

1. Op的具体的Kernel实现，用户通过继承OpKernel类，实现Compute方法，并在其中实现具体Op计算逻辑，还可实现兼容不同设备的实现。

template <typename Device, typename T>
class ZerosLikeOp : public OpKernel {
 public:
  explicit ZerosLikeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}

    void Compute(OpKernelContext* ctx) override {
     ...
    }
  }
};

1. 然后通过REGISTER_KERNEL_BUILDER将Kernel 函数添加到注册表中， 相当于创建一个 OpDefBuilder 对象，并绑定了设备，输入输出等描述。

#define REGISTER_KERNEL(type, dev)                                      
  REGISTER_KERNEL_BUILDER(                                              
      Name("ZerosLike").Device(DEVICE_##dev).TypeConstraint<type>("T"), 
      ZerosLikeOp<dev##Device, type>)

1. 使用REGISTER_OP 宏完成 OpDef 的注册，OpDef 仓库在 C++ 系统 main 函数启动之前完成 OpDef 的加载和注册，后面我们再详细介绍这个过程。

REGISTER_OP("ZerosLike")
    .Input("x: T")
    .Output("y: T")
    .Attr("T: type")
    .SetShapeFn(shape_inference::UnchangedShape);

通过定义REGISTER_OP 系统自动完成字符串表示的翻译表达，并将其转换为 OpDef 的内部表示，最后保存在 OpDef 的仓库中，在 C++ 系统 main 函数启动之前完成加载和注册。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

通过上述的REGISTER_OP，最终会被转换为如下的OpDef ，并最终被保存在OpDef 的仓库，在启动前完成注册。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

op {
        name: "ZerosLike"
        input_arg {
          name: "x"
          type_attr: "T"
        }
        input_arg {
          name: "y"
          type_attr: "T"
        }
        output_arg {
          name: "z"
          type_attr: "T"
        }
        attr {
          name: "T"
          type: "type"
          allowed_values {
            list {
              type: DT_HALF
              type: DT_FLOAT
              type: DT_DOUBLE
              type: DT_UINT8
              type: DT_INT8
              type: DT_UINT16
              type: DT_INT16
              type: DT_INT32
              type: DT_INT64
              type: DT_COMPLEX64
              type: DT_COMPLEX128
            }
          }
        }
        is_commutative: true
      }

每一个Node的都会持有OpDef与NodeDef信息，这样就绑定到具体的Op。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

3. TF Op执行过程

1. 创建计算图：通过TensorFlow的前端Api， 构建一个有向无环计算图（DAG）。TensorFlow 的计算图包含了各种 OP 节点和它们之间的数据流关系，同时也包含了各个 OP 执行所依赖的设备信息、控制流信息、梯度信息等。
2. 计算机的传递与恢复：当前端构造完计算图后，TensorFlow 实际上并不需要再次去构造一遍。而是会将前端构造的计算图序列化成 Protocol Buffer 格式（即 .pb 文件），然后传递给后端运行，后端运行时会根据 .pb 文件中的信息来恢复计算图。
3. 进行图优化与生成可执行计算图：将计算图通过图优化器（Graph Optimizer）进行各种优化，这里主要包括：(1)前端优化：这一部分主要针对前端构造的计算图进行各种优化和简化，以减少不必要的计算和存储开销。例如，TensorFlow 会自动合并多个 OP 节点为一个节点、去除不必要的控制流节点、减少张量传输等。（2）中间表示的优化：在进行前端优化之后，TensorFlow 会将计算图转换为一种称为 GraphDef 或 Grappler Graph 的中间表示形式。然后，TensorFlow 会针对这种中间表示形式进行各种优化和转换，以进一步提高计算性能和效率。例如，TensorFlow 可以应用各种优化算法，如常量折叠、常量传播、死代码消除和算子融合等技术，来尽可能地减少计算和存储开销。 （3）后端优化：最后，TensorFlow 的图优化器还会针对特定的硬件平台进行优化和适配，以进一步提高计算性能和效率。例如，TensorFlow 可以根据硬件平台的具体特性，对计算图进行分割、并行化、向量化、本地化等优化，从而更好地利用硬件资源，提高计算效率和吞吐量。
4. 执行 OP 计算：TensorFlow 会自动将数据传递给各个 OP 节点，由节点的 Kernel 函数进行计算。Kernel 函数会将输入张量转换为 C++ 或 CUDA 数据结构，并进行具体的计算，最终将结果存储在输出张量中。在计算过程中，TensorFlow 还会进行各种优化和调度，以提高计算效率和并行性。
5. 返回输出结果：在 OP 计算完成后，TensorFlow 会将计算结果从输出张量中读取，并返回给用户。用户可以在 Python 中使用 Session 对象的 run() 方法获取输出结果。

4. 总结

本文浅析了TF中Operation（操作）的构成元素和实现，以及它在计算图中的作用。每个 Operation 都代表了计算图中的一个Node节点。它就类似于Spark SQL图中的节点，负责某种抽象计算，相比于Spark，TF拥有更多的Op。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

TF分为前端和后端系统，前端系统负责构图，其中的Operation代表Op节点，其维护了Op所需要的最本质的信息，包括Op的输入输出，属性等，此外也包含了Op间的上游信息。通过Operation的构造器完成Op的创建进而完成图的创建。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html

然后将前端构造的计算图序列化成 Protocol Buffer 格式发送到后端，后端运行时会根据 .pb 文件中的信息来恢复计算图。在后端Operation对应的是Node对象，其内部维护的OpDef则记录具体Op的元信息，在真正运行图计算时，会根据OpDef去，OpDef仓库里找到对应的Op, 并根据设备，选择性能最优的要执行kernel，进行执行。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/43457.html