刘光聪 | TensorFlow：揭示多语言编程的奥秘

最新推荐文章于 2022-09-14 10:43:34 发布

LeadAI学院

最新推荐文章于 2022-09-14 10:43:34 发布

阅读量1.2k

点赞数

北京 | 深度学习与人工智能研修

12月23-24日

再设经典课程重温深度学习

阅读全文

作者：刘光聪，中兴通讯高级系统架构师，专注机器学习算法，分布式系统架构与优化。

正文共2598个字 7张图，预计阅读时间7分钟。

Client是TensorFlow系统中的重要组件。Client负责计算图的构造，并通过Session为桥梁，连接TensorFlow后端的「运行时」，启动计算图的执行过程。

TensorFlow的前端是一个支持多语言的编程环境，支持Client使用多语言构造计算图。本文通过Session创建过程为例，揭示前端Python与后端C/C++系统实现的通道。

系统架构

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

前端系统：提供编程模型，负责构造计算图；
后端系统：提供运行时环境，负责执行计算图。

系统架构

Swig:幕后英雄

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

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的具体实现。

Swig代码生成器

编程接口：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

生成代码

在pywrap_tensorflow模块中，通过_pywrap_tensorflow将在_pywrap_tensorflow.so中调用对应的C++函数实现。

# tensorflow/bazel-bin/tensorflow/python/pywrap_tensorflow.py

def TF_NewDeprecatedSession(arg1, status):

return _pywrap_tensorflow.TF_NewDeprecatedSession(arg1, status)

在pywrap_tensorflow.cpp的具体实现中，它静态注册了函数调用的符号表，实现Python的函数名称到C++实现函数的具体映射。

// tensorflow/bazel-bin/tensorflow/python/pywrap_tensorflow.cpp
static PyMethodDef SwigMethods[] = { ... {"TF_NewDeprecatedSession", _wrap_TF_NewDeprecatedSession, METH_VARARGS, NULL}, } PyObject *_wrap_TF_NewDeprecatedSession( PyObject *self, PyObject *args) { TF_SessionOptions* arg1 = ... TF_Status* arg2 = ... TF_DeprecatedSession* result = TF_NewDeprecatedSession(arg1, arg2);
// ignoring others implements

}

最终，自动生成的pywrap_tensorflow.cpp仅仅负责函数调用的转发，最终将调用底层C系统向上提供的API接口。

CAPI

c_api.h是TensorFlow的后端执行系统面向前端开放的公共API接口之一，自此将进入TensorFlow后端系统的浩瀚天空。

// 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; } }