一文读懂TensorFlow的框架结构

2023-12-10旅游

TensorFlow框架结构是指TensorFlow的系统设计和实现 ，它包括以下几个方面：

TensorFlow基于数据流图（Data Flow Graph） ，用于大规模分布式数值计算的开源框架1\。数据流图由节点（Node）和边（Edge）组成，节点表示计算单元（如算术运算、变量、常量等），边表示数据单元（如张量、变量等）在节点间的传递。

TensorFlow以C API为界，将整个系统分为前端和后端两个子系统 。前端系统提供编程模型，负责构造计算图；后端系统提供运行时环境，负责执行计算图。

TensorFlow的前端系统支持多种语言接口 ，如Python、C++、Java等，用户可以通过这些接口定义计算图和控制其执行。

TensorFlow的后端系统遵循良好的分层架构，可以进一步分解为四层：

运行时（Runtime）： 分别提供本地模式和分布式模式，并共享大部分设计和实现；

计算层（Computation Layer ）：由各个OP的Kernel实现组成，在运行时，Kernel实现执行OP的具体数学运算；

通信层（Communication Layer） ：基于gRPC实现组件间的数据交换，并能够在支持IB网络的节点间实现RDMA通信；

设备层（Device Layer） ：计算设备是OP执行的主要载体，TensorFlow支持多种异构的计算设备类型，如CPU、GPU等。

TensorFlow的模型构造和执行流程一般包括以下几个步骤：

图构建（Graph Construction）：用户在客户端（Client）中基于TensorFlow的多语言编程接口，添加算子（OP），完成计算图的构造。
图传递（Graph Transfer）：客户端开启会话（Session），通过它建立和主控节点（Master）之间的连接。执行Session.run()时，将构造好的Graph序列化为GraphDef后，以Protobuf的格式传递给Master。
图剪枝（Graph Pruning）：Master根据Session.run()传递的Fetches和Feeds列表，反向遍历全图（Full Graph），实施剪枝，得到最小依赖子图（Minimal Dependent Subgraph）。
图分裂（Graph Partitioning）：Master将最小子图分裂为多个图分片（Graph Partition），并注册到多个工作节点（Worker）上。一个Worker对应一个图分片。
图二次分裂（Graph Repartitioning）：Worker根据当前可用硬件资源，如CPU、GPU等，将图分片按照OP算子设备约束规范（例如tf.device(‘/cpu:0’)），二次分裂到不同设备上。每个计算设备对应一个图分片。
图运行（Graph Execution）：对于每一个计算设备，Worker依照OP在Kernel中的实现，完成OP的运算。设备间数据通信可以使用Send/Recv节点，而Worker间通信，则使用gRPC或RDMA协议。