后端session前端访问_Tensorflow源码剖析：session创建

在TensorFlow中&＃xff0c;用户是通过运行图来进行模型训练的&＃xff0c;而启动图的第一步就是创建一个session对象。在日常编写Python代码时&＃xff0c;有的直接通过编写sess&＃61;tf.Session()来创建session&＃xff0c;也有的在分布式TensorFlow中通过ChiefSessionCreator和WorkerSessionCreator的create_session()来创建session。这里简单说明下&＃xff0c;create_session()实质上对tf.Session()的封装&＃xff0c;只是里面添加了很多其他的功能&＃xff0c;后期会对SessionCreator进行详细的介绍。鉴于前期读者反馈说看不大懂&＃xff0c;所以今天&＃xff0c;谱哥主要是想带大家来了解下sess&＃61;tf.Session()背后的实现原理&＃xff0c;并介绍allocator在session创建时在哪里有体现。

 TensorFlow系统分为前端系统和后端系统&＃xff0c;前端系统提供编程模型&＃xff0c;重点负责图的构造&＃xff0c;目前主流编程语言是Python&＃xff1b;后端系统主要负责图的执行&＃xff0c;用C&＃43;&＃43;语言来进行编写&＃xff1b;Swig作为前端系统和后端系统建立连接的桥梁&＃xff0c;使得前端Python创建session能够触发后端C&＃43;&＃43;进行session创建。因此&＃xff0c;接下来&＃xff0c;将按照前端Python层、Swig以及后端C&＃43;&＃43;层三个方面来详细说明sess&＃61;tf.Session()底部实现原理。1. 前端&＃xff1a;Python层在前端系统中&＃xff0c;session相关类的继承关系如下所示&＃xff1a; 

从中可知&＃xff0c;session分为两种&＃xff0c;普通Session和交互式InteractiveSession。后者自带with上下文管理器&＃xff0c;并且在初始化的时候将自身作为默认的session&＃xff0c;因此适合在Python交互式环境下使用。普通Session和交互式InteractiveSession都继承BaseSession&＃xff0c;BaseSession继承SessionInterface。当用户层执行sess&＃61;tf.Session()时&＃xff0c;会依次调用SessionInterface、BaseSession和Session的初始化函数。在BaseSession的初始化函数中有以下几行代码&＃xff1a;

from tensorflow.python import pywrap_tensorflow as tf_sessionclass BaseSession(SessionInterface): def __init__(self, target&＃61;&＃39;&＃39;, graph&＃61;None, config&＃61;None): ...... self._session &＃61; None opts &＃61; tf_session.TF_NewSessionOptions(target&＃61;self._target, config&＃61;config) try: # pylint: disable&＃61;protected-access self._session &＃61; tf_session.TF_NewSession(self._graph._c_graph, opts) # pylint: enable&＃61;protected-access finally: tf_session.TF_DeleteSessionOptions(opts)由上可知&＃xff0c;session是在BaseSession初始化的时候执行tf_session.TF_NewSession()来创建&＃xff0c;传入的参数opts通过tf_session.TF_NewSessionOptions创建&＃xff0c;是一个SessionOptions结构体&＃xff0c;完成env、target和config的简单封装。target参数主要用来判断是创建DirectSession还是GrpcSession。

struct SessionOptions { Env* env; string target; ConfigProto config; SessionOptions();};tf_session实质指的是pywrap_tensorflow.py模块&＃xff0c;该模块内部导入了pywrap_tensorflow_internal.py模块。而pywrap_tensorflow_internal.py是在系统启动Swig的时候通过tensorflow.i自动生成的适配文件&＃xff0c;因此要想知道tf_session.TF_NewSession()内部到底干了啥&＃xff0c;需要了解TensorfFlow是怎样使用Swig。2. Swig包装器在TensorFlow启动Swig的时候&＃xff0c;会通过tensorflow.i生成两个文件&＃xff1a; (1) pywrap_tensorflow_internal.py:对接前端python接口调用&＃xff1b;
(2) pywrap_tensorflow_internal.cc:对接后端C&＃43;&＃43;接口调用。 pywrap_tensorflow_internal.py模块在pywrap_tensorflow.py模块中被导入的时候&＃xff0c;会自动加载_pywrap_tensorflow_internal.so的动态链接库&＃xff0c;该库包含了整个TensorFlow运行时的所有符号。因此&＃xff0c;在pywrap_tensorflow_internal.py模块中&＃xff0c;可以通过_pywrap_tensorflow_internal转发&＃xff0c;实现Python接口到_pywrap_tensorflow_internal.so的函数调用。

def TF_NewSession(graph, opts): return _pywrap_tensorflow_internal.TF_NewSession(graph, opts)TF_NewSession &＃61; _pywrap_tensorflow_internal.TF_NewSessionpywrap_tensorflow_internal.cc注册了一个函数符号表&＃xff0c;实现Python函数到C函数名的映射。

{ (char *)"TF_NewSession", _wrap_TF_NewSession, METH_VARARGS, NULL},而_wrap_TF_NewSession将调用c_api.h对其开放的API接口&＃xff1a;TF_NewSession&＃xff0c;从而进入系统后端C&＃43;&＃43;层。3. 后端&＃xff1a;C&＃43;&＃43;层3.1 TF_NewSession在c_api.cc文件中&＃xff0c;TF_NewSession()相关定义如下所示&＃xff1a;

TF_Session* TF_NewSession(TF_Graph* graph, const TF_SessionOptions* opt, TF_Status* status) { Session* session; //创建session status->status &＃61; NewSession(opt->options, &session); if (status->status.ok()) { //创建TF_Session对象&＃xff0c;实现session和graph的绑定 TF_Session* new_session &＃61; new TF_Session(session, graph); if (graph !&＃61; nullptr) { mutex_lock l(graph->mu); graph->sessions[new_session] &＃61; ""; } //返回TF_Session对象 return new_session; } else { DCHECK_EQ(nullptr, session); return nullptr; }}从上可以看出&＃xff0c;TF_NewSession()干了两件事儿&＃xff1a;(1)创建了session对象&＃xff1b;(2)创建TF_Session对象&＃xff0c;实现session和graph的绑定&＃xff0c;并最终返回TF_Session&＃xff0c;而不是session。 TF_Session的相关定义如下&＃xff1a;

struct TF_Session { TF_Session(tensorflow::Session* s, TF_Graph* g); tensorflow::Session* session; TF_Graph* const graph; int last_num_graph_nodes; // If true, TF_SessionRun and similar methods will call // ExtendSessionGraphHelper before running the graph std::atomic extend_before_run;};

3.2 NewSession

那么NewSession()又是怎么定义的呢&＃xff1f;于是追踪到了session.cc文件&＃xff0c;相关代码如下所示&＃xff1a;

Status NewSession(const SessionOptions& options, Session** out_session) { SessionFactory* factory; const Status s &＃61; SessionFactory::GetFactory(options, &factory); if (!s.ok()) { *out_session &＃61; nullptr; LOG(ERROR) <NewSession(options); if (!*out_session) { return errors::Internal("Failed to create session."); } return Status::OK();}可知&＃xff0c;NewSession采用了工厂模式&＃xff0c;先根据options去找出符合要求的工厂factory&＃xff0c;然后在指定的工厂里创建Session。3.3 SessionFactory::GetFactory 可能大家又会问又是如何去根据options找出对应的factory&＃xff1f;于是&＃xff0c;追踪到了session_factory.cc文件&＃xff0c;在剖析GetFactory()之前&＃xff0c;需要先来理解session_factories()的相关概念&＃xff0c;如下所示&＃xff1a;

typedef std::unordered_map SessionFactories;SessionFactories* session_factories() { static SessionFactories* factories &＃61; new SessionFactories; return factories;}可知&＃xff0c;session_factories()其实是创建了一个静态的SessionFactories&＃xff0c;而这个SessionFactories是一个unordered_map&＃xff0c;实现string类型的runtime_type到SessionFactory指针的映射。那么既然是unordered_map&＃xff0c;就必然涉及到key和value的存储&＃xff0c;在这里是通过SessionFactory::Register()来进行注册的。SessionFactory::Register()的相关定义如下所示&＃xff1a;

void SessionFactory::Register(const string& runtime_type, SessionFactory* factory) { mutex_lock l(*get_session_factory_lock()); if (!session_factories()->insert({runtime_type, factory}).second) { ...... }}由上可知&＃xff0c;SessionFactory::Register()的本质就是将runtime_type和factory指针插入到unordered_map中。紧接着&＃xff0c;我追踪到direct_session.cc和grpc_session.cc文件&＃xff0c;发现了相关session的注册。 DirectSessionFactory的注册如下所示&＃xff1a;

class DirectSessionRegistrar { public: DirectSessionRegistrar() { SessionFactory::Register("DIRECT_SESSION", new DirectSessionFactory()); }};static DirectSessionRegistrar registrar;GrpcSessionFactory的注册如下所示&＃xff1a;

class GrpcSessionRegistrar { public: GrpcSessionRegistrar() { SessionFactory::Register("GRPC_SESSION", new GrpcSessionFactory()); }};static GrpcSessionRegistrar registrar;现在来看看SessionFactory::GetFactory()核心代码&＃xff0c;如下所示&＃xff1a;

Status SessionFactory::GetFactory(const SessionOptions& options, SessionFactory** out_factory) { std::vector<:pair sessionfactory>> candidate_factories; for (const auto& session_factory : *session_factories()) { if (session_factory.second->AcceptsOptions(options)) { candidate_factories.push_back(session_factory); } } if (candidate_factories.size() &＃61;&＃61; 1) { *out_factory &＃61; candidate_factories[0].second; return Status::OK(); } else if (candidate_factories.size() > 1) { //报错 }}从GetFactory()代码中可知&＃xff0c;其本质就是遍历session_factories()中的unordered_map&＃xff0c;然后通过unordered_map中的SessionFactory(如DirectSessionFactory和GrpcSessionFactory)是否AcceptsOptions来进行选择&＃xff0c;并且硬性要求有且仅有一个factory满足要求&＃xff0c;否则报错。不同SessionFactory的AcceptsOptions()的定义如下&＃xff1a; DirectSessionFactory:

bool AcceptsOptions(const SessionOptions& options) override { return options.target.empty(); }从上可知&＃xff0c;若options.target为空&＃xff0c;则应选择DirectSessionFactory,用于本地训练。 GrpcSessionFactory:

const char* const kSchemePrefix &＃61; "grpc://";bool AcceptsOptions(const SessionOptions& options) override { return str_util::StartsWith(options.target, kSchemePrefix); }从上可知&＃xff0c;若options.target是以grpc://开头的&＃xff0c;则应选择GrpcSessionFactory&＃xff0c;用于分布式TensorFlow。3.4 factory->NewSession延续3.2节所讲&＃xff0c;根据3.3节SessionFactory::GetFactory返回值的SessionFactory类型&＃xff0c;去调用对应SessionFactory的NewSession接口。这里以DirectSessionFactory::NewSession为例&＃xff0c;代码如下&＃xff1a;

Session* NewSession(const SessionOptions& options) override { // Must do this before the CPU allocator is created. if (options.config.graph_options().build_cost_model() > 0) { EnableCPUAllocatorFullStats(true); } std::vector devices; const Status s &＃61; DeviceFactory::AddDevices( options, "/job:localhost/replica:0/task:0", &devices); if (!s.ok()) { LOG(ERROR) <从上可知&＃xff0c;DirectSessionFactory::NewSession()不单单只是创建DirectSession&＃xff0c;还要完成相关的devices收集。相关设备收集通过调用 DeviceFactory::AddDevices()来完成&＃xff0c;相关代码在device_factory.cc中&＃xff0c;如下所示&＃xff1a;

Status DeviceFactory::AddDevices(const SessionOptions& options, const string& name_prefix, std::vector* devices) { //先获取CPU对应的设备工厂cpu_factory auto cpu_factory &＃61; GetFactory("CPU"); //创建设备并记录保存到devices TF_RETURN_IF_ERROR(cpu_factory->CreateDevices(options, name_prefix, devices)); ... //遍历device_factories()&＃xff0c;创建设备集&＃xff0c;包括GPU for (auto& p : device_factories()) { auto factory &＃61; p.second.factory.get(); if (factory !&＃61; cpu_factory) { TF_RETURN_IF_ERROR(factory->CreateDevices(options, name_prefix, devices)); } } return Status::OK();}

从上可知&＃xff0c;DeviceFactory::AddDevices()也是采用工厂模式&＃xff0c;主要完成的是遍历device_factories()&＃xff0c;然后调用每个factory中的CreateDevices接口&＃xff0c;创建设备并把相应指针存储到devices vector中。在此&＃xff0c;有几个接口函数需要说明下&＃xff1a;

DeviceFactory::GetFactory

DeviceFactory* DeviceFactory::GetFactory(const string& device_type) { //根据device_type查找对应的DeviceFactory auto it &＃61; device_factories().find(device_type); if (it &＃61;&＃61; device_factories().end()) { return nullptr; } return it->second.factory.get();}DeviceFactory::GetFactory()主要是通过输入参数device_type在device_factories()中查找对应的设备工厂DeviceFactory。device_factories()的本质类似于session_factories()&＃xff0c;函数里创建了一个静态的unordered_map&＃xff0c;表示device_type到FactoryItem的映射。FactoryItem是个结构体&＃xff0c;包括factory指针和相应的优先级。

struct FactoryItem { std::unique_ptr factory; int priority;};std::unordered_map& device_factories() { static std::unordered_map* factories &＃61; new std::unordered_map; return *factories;}跟SessionFactory::Register()类似&＃xff0c;既然涉及到对unordered_map的读取&＃xff0c;那么肯定存在对key和value的存储操作。该操作主要是通过DeviceFactory::Register()接口来完成相关DeviceFactory的注册。在TensorFlow中&＃xff0c;专门为此进行了宏定义&＃xff0c;如下所示&＃xff1a;

#define REGISTER_LOCAL_DEVICE_FACTORY(device_type, device_factory, ...)INTERNAL_REGISTER_LOCAL_DEVICE_FACTORY(device_type, device_factory, __COUNTER__, ##__VA_ARGS__)以下是相关DeviceFactory注册的部分代码&＃xff0c;分别在threadpool_device_factory.cc和gpu_device_factory.cc文件中。从这也可以看出&＃xff0c;GPUDeviceFactory的优先级要明显高于ThreadPoolDeviceFactory。

REGISTER_LOCAL_DEVICE_FACTORY("CPU", ThreadPoolDeviceFactory, 60);REGISTER_LOCAL_DEVICE_FACTORY("CPU", GPUCompatibleCPUDeviceFactory, 70);REGISTER_LOCAL_DEVICE_FACTORY("GPU", GPUDeviceFactory, 210);CreateDevices 遍历device_factories()里unordered_map的时候&＃xff0c;都会让每个DeviceFactory调用设备创建CreateDevices()&＃xff0c;并存储到std::vector* devices中。下面以ThreadPoolDeviceFactory::CreateDevices为例来介绍其具体细节。

Status ThreadPoolDeviceFactory::CreateDevices(const SessionOptions& options, const string& name_prefix, std::vector* devices) override { int n &＃61; 1; auto iter &＃61; options.config.device_count().find("CPU"); if (iter !&＃61; options.config.device_count().end()) { n &＃61; iter->second; } for (int i &＃61; 0; i push_back(new ThreadPoolDevice( options, name, Bytes(256 <<20), DeviceLocality(), cpu_allocator())); } return Status::OK(); }从第9行可以看出&＃xff0c;ThreadPoolDeviceFactory::CreateDevices主要是创建了一个ThreadPoolDevice&＃xff0c;ThreadPoolDevice中存有一个allocator用来分配和释放内存。该allocator由cpu_allocator()来获取。cpu_allocator()相关定义在allocator.cc文件中&＃xff0c;如下所示&＃xff1a;

Allocator* cpu_allocator() { static Allocator* cpu_alloc &＃61; AllocatorRegistry::Global()->GetAllocator(); if (cpu_allocator_collect_full_stats && !cpu_alloc->TracksAllocationSizes()) { cpu_alloc &＃61; new TrackingAllocator(cpu_alloc, true); } return cpu_alloc;}不知道大家现在对第2行的接口有没有丝丝熟悉感&＃xff1f;对的&＃xff0c;在Allocator(基础篇)已经对AllocatorRegistry相关接口进行了说明。从第2行可看出&＃xff0c;每次执行AllocatorRegistry::Global()->GetAllocator()都会返回AllocatorRegistry当前优先级最高的allocator。如果该allocator想收集状态但是TracksAllocationSizes()又为false&＃xff0c;那么就可以对该allocator进行封装&＃xff0c;在此基础上创建一个TrackingAllocator即可进行记录追踪。当然&＃xff0c;如果想知道底层是在哪里使用了BFCAllocator&＃xff0c;则推荐阅读GPUDeviceFactory::CreateDevices&＃xff0c;这里不再做过多说明。4. 总结本篇根据前端Python层、Swig以及后端C&＃43;&＃43;层三个方面来详细说明sess&＃61;tf.Session()底部实现原理。前端Python层介绍了Session和BaseSession等的概念和相互联系&＃xff1b;Swig主要完成将Python层的session创建转发到C&＃43;&＃43;层的session创建&＃xff1b;后端C&＃43;&＃43;层session创建根据SessionOptions找到相应的SessionFactory来执行NewSession操作。而NewSession函数不仅要完成session创建&＃xff0c;还要根据DeviceFactory完成设备的创建并收集&＃xff0c;这里自然离不开各种allocator。