探索TensorFlow核心组件系列的基本机制

terry 2年前 (2023-09-23) 阅读数 72 #AI人工智能

在TensorFlow中，Operation（运算）是核心组件之一，就是计算图。它代表计算图中的节点并执行各种数学运算、数据处理和转换操作。

其中，每个Operation都有自己的计算逻辑和属性来体现他所做的具体操作。运算范围可以从简单的数学运算（例如加法和乘法）到复杂的神经网络层或专用算法。

Operation的基本机制包括以下主要元素：计算函数、输入输出张量、属性、内存管理。这些元素共同作用，使Operation能够有效地执行计算任务并产生结果。可以说，Operation的基本机制是TensorFlow高效计算和灵活模型构建的关键。

TensorFlow 中的操作类似于 Spark 中的计算节点（或节点）。它负责某种抽象计算。这代表计算图中的一个节点。

但相比Spark中Node计算节点的分类，TF更为复杂，主要包括：

OP数学运算，主要进行加、减、乘、除、矩阵乘、矩阵等运算。转置；

OP神经网络，主要做神经网络建模，包括卷积、池化、全连接等；

OP数据操作，特别是处理输入输出数据；

OP Optimizer，主要在训练过程中进行模型优化操作，如梯度下降、自适应学习率优化算法；

OP流控制，主要用于控制程序运行过程，如条件语句、循环语句等；

还有图像运算、分布公式、数据集、Op。

我们都知道，在TensorFlow中，前端负责组合，后端负责运行，那么我们分别从前端和后端看一下Op的实现逻辑。

1。前端Op的实现（Python）

1.1 创建图从创建Op开始

在计算图的构建过程中，通过Op构造函数创建了一个Operation的实例，在创建过程中注册到默认图表。

之间，类元数据Operation 由 OpDef 和 NodeDef 保存。它们是ProtoBuf格式，描述了关于Operation最重要、最重要的事情。其中，OpDef描述了OP的静态属性信息，如OP名称、输入/输出参数列表、属性集定义等信息。 NodeDef表示OP的动态属性值信息，如属性值等信息。我们稍后会详细解释。

1.2 通过Operation构造函数创建Op

在Python前端，Op对象是通过Operation构造函数创建的。具体代码如下：

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相关。通过这个 Dependency 可以在 session.run() 打开时在最小的依赖子图中找到。

从Operation类的成员变量可以看出，它主要包含两个主要的成员变量：OpDef和NodeDef。它还列出了图形信息、输入和输出类型、输入张量列表和控制输入列表。并通过遍历所有上游输入张量并将当前 Operation 添加到维护的消费者列表中来创建生产者-消费者关系。

我们看一下NodeDef的组成：

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包含了Op的名称、其属性以及当前Node的输入等信息。

需要注意的是，Operation的一些创作是不需要输入的，比如Constant。此类操作称为资源操作。

1.3 创建的Op对象作为Node节点添加到Graph中，最后创建Session并执行Op。

将图从OP倒置到末尾，根据依赖关系找到依赖关系最小的子图。寻找依赖的过程就是寻找输入，直到没有输入，并执行默认Session中的子图。

注意，第一次执行图表时，会分配一个Executor。但是，如果图没有改变并且再次支持，则最后分配的Executor将立即返回。

在前端系统中，Operation创建一个Op来构建整个Graph，然后将其序列化为Protocol Buffer格式（即.pb文件）并转发到后端系统。。

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

在C++后端系统中，Operation的Python前端对应后端Node实现。

节点（node）可以有零个或多个输入/输出边，并使用 in_edges 和 out_edges 来表示输入边和输出边的集合。另外，后端中的Node对象还保存有NodeDef、OpDef。其中，NodeDef包含设备分配信息和OP属性值列表； OpDef包含OP元数据，包括OP输入输出类型等信息。节点和边一起形成一个图。

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等组成信息。

在大家心中，Op代表的是计算函数，比如tf.divide，引入当前图中的Node对象。

那么 Node 对象如何与特定 Op 相关联？

其实主要依赖于Node中的OpDef对象。下面详细讲一下Op组件的定义。

TensorFlow中的所有Op定义主要包括以下部分。下面我们以 ZerosLikeOp 为例：

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 {
     ...
    }
  }
};

然后通过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>)

使用REGISTER_OP宏完成OpDef注册。 OpDef存储库在主要C++系统功能启动之前完成OpDef加载和注册。我们稍后会详细介绍这个过程。

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

通过指定REGISTER_OP，系统自动从字符串表示中解析翻译表达式，将其转换为内部OpDef表示，最后将其存储在OpDef存储库中。它在主要 C++ 系统函数启动之前加载并注册。

通过上面的REGISTER_OP，最终会转换成下面的OpDef，最后存储在OpDef存储库中，在启动之前完成注册。

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。

3。 TF Op 执行流程

创建计算图：通过 TensorFlow 前端 API 创建有向非循环计算图（DAG）。 TensorFlow计算图包含各种OP节点以及它们之间的数据流关系，以及每个OP执行所依赖的设备信息、控制流信息、梯度信息等。
计算机传输和恢复：前端创建计算图后，TensorFlow不需要重建它。但前端创建的计算图会被放入Protocol Buffer格式（即.pb文件）中，然后转发到后端运行。后端运行时会基于文件.pb。信息来恢复计算图。
进行图优化，创建可执行的计算图：通过图优化器（Graph Optimizer）对计算图进行各种优化，主要包括：（1）前端优化：❀主要针对前端内置的计算图经过了各种优化和简化，以减少不必要的计算和存储开销。例如，TensorFlow会自动将多个OP节点合并为一个节点，去除不必要的控制流节点，减少张量传输等。图计算为称为 GraphDef 或 Grappler Graph 的中间表示。然后，TensorFlow 对此中间表示执行各种优化和转换，以提高性能和计算效率。例如，TensorFlow可以实现各种优化算法，例如恒定折叠、恒定扩展、死代码消除和算子融合，以减少计算和存储开销。 （3）后端优化：最后，TensorFlow图优化器也会针对特定的硬件平台进行优化适配，以提升性能和计算效率。例如，TensorFlow可以根据硬件平台的具体特点，对计算图进行分区、并行化、向量化、局部化等优化，从而更好地利用硬件资源，提高计算效率和吞吐量。
执行OP计算：TensorFlow会自动发送数据到每个OP节点，由节点的Kernel函数执行计算。内核函数将输入张量转换为 C++ 或 CUDA 数据结构，执行某些计算，并将结果存储在输出张量中。在计算过程中，TensorFlow还会进行各种优化和调度，以提高计算效率和并行性。
返回输出结果：OP计算完成后，TensorFlow会从输出张量中读取计算结果并返回给用户。用户可以使用Python中Session对象的run()方法来获取输出。

4。总结

本文简单分析了TF中Operation（操作）的构成要素和实现，以及它们在计算图中的作用。每个Operation代表计算图中的一个Node节点。它类似于Spark SQL图中的节点，负责一些抽象计算。相比Spark，TF有更多的Ops。

TF分为前端系统和后端系统。前端系统负责创建图像。 Operation代表Op节点，它维护着Op所需要的最重要的信息，包括Op的输入、输出、属性等。此外，还包括节点Op。上游信息。通过 Operation 构造函数完成 Op 的创建，然后完成图形的创建。

然后将前端创建的计算图输入Protocol Buffer格式并发送给后端。后端运行时，会根据文件.pb中的信息恢复计算图。在后端，Operation对应Node对象，内部维护OpDef记录特定Op的元信息。实际运行图计算时，会根据OpDef在OpDef存储库中查找合适的Op，并根据设备选择性能最佳的Op。最好的办法是运行内核并运行它。