TensorFlow 学习之二基础架构

TensorFlow支持各种异构的平台，支持多CPU/GPU，服务器，移动设备，具有良好的跨平台的特性；
TensorFlow的系统结构以C API为界，将整个系统分为「前端」和「后端」两个子系统：
整的架构图如下：

分层介绍

• 底层： 设备通信层， 负责网络通信和设备管理。 设备管理可以实现TF设备异构的特性，支持CPU、GPU、Mobile等不同设备。网络通信依赖gRPC通信协议实现不同设备间的数据传输和更新。
• 第二层：Kernel 实现层， 以Tensor为处理对象，依赖网络通信和设备内存分配，实现了各种Tensor操作或计算。Opkernels不仅包含MatMul等计算操作，还包含Queue等非计算操作。
• 第三层： 图计算层，包含本地计算流图和分布式计算流图的实现。Graph模块包含Graph的创建、编译、优化和执行等部分。
• 第四层: API接口层， C API是对TF功能模块的接口封装，便于其他语言平台调用。
• 第五层: Client 层，不同编程语言在应用层通过API接口层调用TF核心功能实现相关实验和应用。

核心组件介绍

Client

Client 提供构造computational graph的环境。
之后通过Session连接computational graph后端的计算环境，并启动computational graph的执行过程。

Distributed Master

Distributed Master 主要完成以下几方面的工作：

• 根据Session.run的Fetching参数，从中computational graph反向遍历，找到所依赖的「最小子图」
• 将第一步的「最小子图」分解成子图片段；以便其在不同的进程和设备上执行；
• 将子图片段派发给Work Service； 随后Work Service启动执行这些子图片段的执行过程；

Work Service

Work Service 下面三项工作：

• 处理来自Master的请求；
• 调用Operation的Kernel完成Operation的计算；
• 将Operation的计算结果发送给其它Work Service； 或者接受其它Work Service 发送给它的Operation 的计算结果。

程序运行生命周期

Swig: 幕后英雄

TensorFlow使用Bazel的构建工具，在编译之前启动Swig的代码生成过程，通过tf_session.i自动生成了两个适配(Wrapper)文件：

• pywrap_tensorflow.py: 负责对接上层Python调用；

• pywrap_tensorflow.cpp: 负责对接下层C实现。
  此外，pywrap_tensorflow.py模块首次被加载时，自动地加载_pywrap_tensorflow.so的动态链接库。从而实现了pywrap_tensorflow.py到pywrap_tensorflow.cpp的函数调用关系。
  在pywrap_tensorflow.cpp的实现中，静态注册了一个函数符号表。在运行时，按照Python的函数名称，匹配找到对应的C函数实现，最终转调到c_api.c的具体实现。

  
  

编程接口：Python

当Client要启动计算图的执行过程时，先创建了一个Session实例，进而调用父类BaseSession的构造函数。

# tensorflow/python/client/session.py
class Session(BaseSession):
  def __init__(self, target='', graph=None, config=None):
    super(Session, self).__init__(target, graph, config=config)
    # ignoring others

在BaseSession的构造函数中，将调用pywrap_tensorflow模块中的函数。其中，pywrap_tensorflow模块自动由Swig生成。

# tensorflow/python/client/session.py
from tensorflow.python import pywrap_tensorflow as tf_session
class BaseSession(SessionInterface):
def __init__(self, target='', graph=None, config=None):
self._session = None
opts = tf_session.TF_NewSessionOptions(target=self._target, config=config)
try:
    with errors.raise_exception_on_not_ok_status() as status:
        self._session = tf_session.TF_NewDeprecatedSession(opts, status)
finally:
    tf_session.TF_DeleteSessionOptions(opts)
# ignoring others

生成代码：Swig

pywrap_tensorflow.py将在_pywrap_tensorflow.so中调用对应的C++函数实现。
在pywrap_tensorflow.cpp的具体实现中，它静态注册了函数调用的符号表，实现Python的函数名称到C++实现函数的具体映射。
最终，自动生成的pywrap_tensorflow.cpp仅仅负责函数调用的转发，最终将调用底层C系统向上提供的API接口。

C API：桥梁

c_api.h是TensorFlow的后端执行系统面向前端开放的公共API接口之一

// tensorflow/c/c_api.c
TF_DeprecatedSession* TF_NewDeprecatedSession(
  const TF_SessionOptions*, TF_Status* status) {
  Session* session;
  status->status = NewSession(opt->options, &session);
  if (status->status.ok()) {
    return new TF_DeprecatedSession({session});
  } else {
    return NULL;
  }
}

后端系统：C++

NewSession将根据前端传递的Session.target，使用SessionFactory多态创建不同类型的Session(C++)对象。

Status NewSession(const SessionOptions& options, Session** out_session) {
  SessionFactory* factory;
  Status s = SessionFactory::GetFactory(options, &factory);
  if (!s.ok()) {
    *out_session = nullptr;
    LOG(ERROR) << s;
    return s;
  }
  *out_session = factory->NewSession(options);
  if (!*out_session) {
    return errors::Internal("Failed to create session.");
  }
  return Status::OK();
}

具体会话生命周期

创建会话

上文介绍了Session创建的详细过程，从Python前端为起点，通过Swig自动生成的Python-C++的包装器为媒介，实现了Python到TensorFlow的C API的调用。
其中，C API是前端系统与后端系统的分水岭。后端C++系统根据前端传递的Session.target，使用SessionFactory多态创建Session(C++)对象。

从严格的角色意义上划分，GrpcSession依然扮演了Client的角色。它使用target，通过RPC协议与Master建立通信连接，因此，GrpcSession同时扮演了RPC Client的角色

创建/扩展图

随后，Python前端将调用Session.run接口，将构造好的计算图，以GraphDef的形式发送给C++后端。
其中，前端每次调用Session.run接口时，都会试图将新增节点的计算图发送给后端系统，以便后端系统将新增节点的计算图Extend到原来的计算图中。特殊地，在首次调用Session.run时，将发送整个计算图给后端系统。
后端系统首次调用Session.Extend时，转调(或等价)Session.Create；以后，后端系统每次调用Session.Extend时将真正执行Extend的语义，将新增的计算图的节点追加至原来的计算图中。
随后，后端将启动计算图执行的准备工作。

迭代运行

接着，Python前端Session.run实现将Feed, Fetch列表准备好，传递给后端系统。后端系统调用Session.Run接口。
后端系统的一次Session.Run执行常常被称为一次Step，Step的执行过程是TensorFlow运行时的核心。
每次Step，计算图将正向计算网络的输出，反向传递梯度，并完成一次训练参数的更新。首先，后端系统根据Feed, Fetch，对计算图(常称为Full Graph)进行剪枝，得到一个最小依赖的计算子图(常称为Client Graph)。
然后，运行时启动设备分配算法，如果节点之间的边横跨设备，则将该边分裂，插入相应的Send与Recv节点，实现跨设备节点的通信机制。
随后，将分裂出来的子图片段(常称为Partition Graph)注册到相应的设备上，并在本地设备上启动子图片段的执行过程。

关闭会话

当计算图执行完毕后，需要关闭Session，以便释放后端的系统资源，包括队列，IO等

销毁会话

最后，会话关闭之后，Python前端系统启动GC，当Session.del被调用后，启动后台C++的Session对象销毁过程

总结

本文主要讲解了tensorflow的架构和设计， 罗列了重要的核心组件，以及简单的探讨了程序的生命周期。关于更具体的比如compuatitional graph的生成过程，不同子图片段的分配算法等还没有涉及到。大家可以从文献中找到相应的讲解。

Reference

• http://www.jianshu.com/p/a5574ebcdeab
• http://www.uml.org.cn/zjjs/201704123.asp
• https://zhuanlan.zhihu.com/p/25646408
• http://www.algorithmdog.com/dynamic-tensorflow
• Large Scale Deep Learning with TensorFlow
• Tensorflow: Large-scale machine learning on heterogeneous distributed systems