TensorFlow架构与设计:概述

organicmeercat 发布于2月前 阅读1053次
0 条评论

TensorFlow架构与设计:概述

TensorFlow是什么?

TensorFlow 基于数据流图,用于大规模分布式数值计算的开源框架。节点表示某种抽象的计算,边表示节点之间相互联系的张量。

TensorFlow架构与设计:概述

计算图实例

TensorFlow 支持各种异构的平台,支持多 CPU/GPU ,服务器,移动设备,具有良好的跨平台的特性; TensorFlow 架构灵活,能够支持各种网络模型,具有良好的通用性;此外, TensorFlow 架构具有良好的可扩展性,对 OP 的扩展支持, Kernel 特化方面表现出众。

TensorFlow 最初由 Google 大脑的研究员和工程师开发出来,用于机器学习和神经网络方面的研究,于 2015.10 宣布开源,在众多深度学习框架中脱颖而出,在 Github 上获得了最多的 Star 量。

本文将阐述 TensorFlow 的系统架构,帮助读者加深理解 TensorFlow 的工作机理。

本文假设读者已经了解 TensorFlow 的基本编程模型,包括计算图, OP , Tensor , Session 等基本概念。

系统概述

TensorFlow 的系统结构以 C API 为界,将整个系统分为「前端」和「后端」两个子系统:

  • 前端系统:提供编程模型,负责构造计算图;
  • 后端系统:提供运行时环境,负责执行计算图。

TensorFlow架构与设计:概述

TensorFlow系统架构

如上图所示,重点关注系统中如下4个基本组件,它们是系统分布式运行机制的核心。

Client

Client 是前端系统的主要组成部分,它是一个支持多语言的编程环境。它提供基于计算图的编程模型,方便用户构造各种复杂的计算图,实现各种形式的模型设计。

Client 通过 Session 为桥梁,连接 TensorFlow 后端的「运行时」,并启动计算图的执行过程。

Distributed Master

在分布式的运行时环境中, Distributed Master 根据 Session.run 的 Fetching 参数,从计算图中反向遍历,找到所依赖的「最小子图」。

然后, Distributed Master 负责将该「子图」再次分裂为多个「子图片段」,以便在不同的进程和设备上运行这些「子图片段」。

最后, Distributed Master 将这些「子图片段」派发给 Work Service ;随后 Work Service 启动「子图片段」的执行过程。

Worker Service

对于每以个任务, TensorFlow 都将启动一个 Worker Service 。 Worker Service 将按照计算图中节点之间的依赖关系,根据当前的可用的硬件环境(GPU/CPU),调用 OP 的 Kernel 实现完成 OP 的运算(一种典型的多态实现技术)。

另外, Worker Service 还要负责将 OP 运算的结果发送到其他的 Work Service ;或者接受来自其他 Worker Service 发送给它的 OP 运算的结果。

Kernel Implements

Kernel 是 OP 在某种硬件设备的特定实现,它负责执行 OP 的运算。

组件交互

TensorFlow架构与设计:概述

组件交互

如上图所示,假设存在两个任务:

  • /job:ps/task:0 : 负责模型参数的存储和更新
  • /job:worker/task:0 : 负责模型的训练或推理

接下来,我们将进一步抽丝剥茧,逐渐挖掘出 TensorFlow 计算图的运行机制。

客户端

Client 基于 TensorFlow 的编程接口,构造计算图。目前, TensorFlow 主流支持 Python 和 C++ 的编程接口,并对其他编程语言接口的支持日益完善。

此时, TensorFlow 并未执行任何计算。直至建立 Session 会话,并以 Session 为桥梁,建立 Client 与后端运行时的通道,将 Protobuf 格式的 GraphDef 发送至 Distributed Master 。

也就是说,当 Client 对 OP 结果进行求值时,将触发 Distributed Master 的计算图的执行过程。

如下图所示, Client 构建了一个简单计算图。它首先将 w 与 x 进行矩阵相乘,再与截距 b 按位相加,最后更新至 s 。

TensorFlow架构与设计:概述

构造计算图

Distributed Master

在分布式的运行时环境中, Distributed Master 根据 Session.run 的 Fetching 参数,从计算图中反向遍历,找到所依赖的最小子图。

然后 Distributed Master 负责将该子图再次分裂为多个「子图片段」,以便在不同的进程和设备上运行这些「子图片段」。

最后, Distributed Master 将这些图片段派发给 Work Service 。随后 Work Service 启动「本地子图」的执行过程。

Distributed Master 将会缓存「子图片段」,以便后续执行过程重复使用这些「子图片段」,避免重复计算。

TensorFlow架构与设计:概述

执行图计算

如上图所示, Distributed Master 开始执行计算子图。在执行之前, Distributed Master 会实施一系列优化技术,例如「公共表达式消除」,「常量折叠」等。随后, Distributed Master 负责任务集的协同,执行优化后的计算子图。

子图片段

TensorFlow架构与设计:概述

子图片段

如上图所示,存在一种合理的「子图片段」划分算法。 Distributed Master 将模型参数相关的 OP 进行分组,并放置在 PS 任务上。其他 OP 则划分为另外一组,放置在 Worker 任务上执行。

SEND/RECV节点

TensorFlow架构与设计:概述

插入SEND/RECV节点

如上图所示,如果计算图的边被任务节点分割, Distributed Master 将负责将该边进行分裂,在两个分布式任务之间插入 SEND 和 RECV 节点,实现数据的传递。

随后, Distributed Master 将「子图片段」派发给相应的任务中执行,在 Worker Service 成为「本地子图」,它负责执行该子图的上的 OP 。

Worker Service

对于每个任务,都将存在相应的 Worker Service ,它主要负责如下3个方面的职责:

  • 处理来自 Master 的请求;
  • 调度 OP 的 Kernel 实现,执行本地子图;
  • 协同任务之间的数据通信。

TensorFlow架构与设计:概述

执行本地子图

Worker Service 派发 OP 到本地设备,执行 Kernel 的特定实现。它将尽最大可能地利用多 CPU/GPU 的处理能力,并发地执行 Kernel 实现。

另外, TensorFlow 根据设备类型,对于设备间的 SEND/RECV 节点进行特化实现:

  • 使用 cudaMemcpyAsync 的API实现本地 CPU 与 GPU 设备的数据传输;
  • 对于本地的 GPU 之间则使用端到端的 DMA ,避免了跨 host CPU 昂贵的拷贝过程。

对于任务之间的数据传递, TensorFlow 支持多协议,主要包括:

  • gRPC over TCP
  • RDMA over Converged Ethernet

Kernel Implements

TensorFlow 的运行时包含 200 多个标准的 OP ,包括数值计算,多维数组操作,控制流,状态管理等。每一个 OP 根据设备类型都会存在一个优化了的 Kernel 实现。在运行时,运行时根据本地设备的类型,为 OP 选择特定的 Kernel 实现,完成该 OP 的计算。

TensorFlow架构与设计:概述

TensorFlow Core

其中,大多数 Kernel 基于 Eigen::Tensor 实现。 Eigen::Tensor 是一个使用 C++ 模板技术,为多核 CPU/GPU 生成高效的并发代码。但是, TensorFlow 也可以灵活地直接使用 cuDNN 实现更高效的 Kernel 。

此外, TensorFlow 实现了矢量化技术,使得在移动设备,及其满足高吞吐量,以数据为中心的应用需求,实现更高效的推理。

如果对于复合 OP 的子计算过程很难表示,或执行效率低下, TensorFlow 甚至支持更高效的 Kernle 实现的注册,其扩展性表现相当优越。

技术栈

最后,按照 TensorFlow 的软件层次,通过一张表格罗列 TensorFlow 的技术栈,以便更清晰地对上述内容做一个简单回顾。

TensorFlow架构与设计:概述

TensorFlow技术栈

查看原文: TensorFlow架构与设计:概述

共收到0条回复

需要 登录 后回复方可回复, 如果你还没有账号你可以 注册 一个帐号。