Tensorflow源码解析2 -- 前后端连接的桥梁 - Session

1 Session概述

Session是TensorFlow前后端连接的桥梁。用户利用session使得client能够与master的执行引擎建立连接，并通过session.run()来触发一次计算。它建立了一套上下文环境，封装了operation计算以及tensor求值的环境。

session创建时，系统会分配一些资源，比如graph引用、要连接的计算引擎的名称等。故计算完毕后，需要使用session.close()关闭session，避免引起内存泄漏，特别是graph无法释放的问题。可以显式调用session.close(),或利用with上下文管理器，或者直接使用InteractiveSession。

session之间采用共享graph的方式来提高运行效率。一个session只能运行一个graph实例，但一个graph可以运行在多个session中。一般情况下，创建session时如果不指定Graph实例，则会使用系统默认Graph。常见情况下，我们都是使用一个graph，即默认graph。当session创建时，不会重新创建graph实例，而是默认graph引用计数加1。当session close时，引用计数减1。只有引用计数为0时，graph才会被回收。这种graph共享的方式，大大减少了graph创建和回收的资源消耗，优化了TensorFlow运行效率。
 

2 默认session

op运算和tensor求值时，如果没有指定运行在哪个session中，则会运行在默认session中。通过session.as_default()可以将自己设置为默认session。但个人建议最好还是通过session.run(operator)和session.run(tensor)来进行op运算和tensor求值。

operation.run()

operation.run()等价于tf.get_default_session().run(operation)

@tf_export("Operation")
class Operation(object):
    # 通过operation.run()调用，进行operation计算
    def run(self, feed_dict=None, session=None):
        _run_using_default_session(self, feed_dict, self.graph, session)
        
    def _run_using_default_session(operation, feed_dict, graph, session=None):
          # 没有指定session，则获取默认session
          if session is None:
            session = get_default_session()
            
         # 最终还是通过session.run()进行运行的。tf中任何运算，都是通过session来run的。
          # 通过session来建立client和master的连接，并将graph发送给master，master再进行执行
           session.run(operation, feed_dict)

tensor.eval()

tensor.eval()等价于tf.get_default_session().run(tensor), 如下

@tf_export("Tensor")
class Tensor(_TensorLike):
    # 通过tensor.eval()调用，进行tensor运算
    def eval(self, feed_dict=None, session=None):
        return _eval_using_default_session(self, feed_dict, self.graph, session)
        
    def _eval_using_default_session(tensors, feed_dict, graph, session=None):
           # 如果没有指定session，则获取默认session
          if session is None:
            session = get_default_session()
            return session.run(tensors, feed_dict)

默认session的管理

tf通过运行时维护的session本地线程栈，来管理默认session。故不同的线程会有不同的默认session，默认session是线程作用域的。

# session栈
_default_session_stack = _DefaultStack()

# 获取默认session的接口
@tf_export("get_default_session")
def get_default_session():
  return _default_session_stack.get_default()

# _DefaultStack默认session栈是线程相关的
class _DefaultStack(threading.local):
      # 默认session栈的创建，其实就是一个list
      def __init__(self):
        super(_DefaultStack, self).__init__()
        self._enforce_nesting = True
        self.stack = []  
        
      # 获取默认session
      def get_default(self):
        return self.stack[-1] if len(self.stack) >= 1 else None

 

3 前端Session类型

session类图

会话Session的UML类图如下

分为两种类型，普通Session和交互式InteractiveSession。InteractiveSession和Session基本相同，区别在于

1. InteractiveSession创建后，会将自己替换为默认session。使得之后operation.run()和tensor.eval()的执行通过这个默认session来进行。特别适合Python交互式环境。
2. InteractiveSession自带with上下文管理器。它在创建时和关闭时会调用上下文管理器的enter和exit方法，从而进行资源的申请和释放，避免内存泄漏问题。这同样很适合Python交互式环境。

Session和InteractiveSession的代码逻辑不多，主要逻辑均在其父类BaseSession中。主要代码如下

@tf_export('Session')
class Session(BaseSession):
    def __init__(self, target='', graph=None, config=None):
        # session创建的主要逻辑都在其父类BaseSession中
        super(Session, self).__init__(target, graph, config=con