目录

  1. 核心概念
  2. direct_session
    1. direct_session.h
    2. direct_session.cc

1. 核心概念

读过之前文章的读者应该还记得,session是一个执行代理。我们把计算图和输入交给session,由它来调度执行器,执行计算产生结果。TF给我们提供了一个最简单的执行器direction_session。按照当前的理解,我们觉得direction_session的实现应该是非常简单而直接的,毕竟执行器的复杂结构我们在executor那篇已经见到了。但实际上,问题的难点在于,有时候我们只是希望以计算图中某些节点为输入,某些节点为输出,来执行图中的一小部分计算,而不需要执行整张图,另外一个方面,这种对图部分执行的任务,在同一张图上可能同时存在多个。为了应对这种情况,direct_session就衍生出了很多辅助数据。

2. direct_session

2.1 direct_session.h

DirectSession类提供了丰富的数据和接口,以下为了表达简洁,我们略去了部分函数的形参:

class DirectSession : public Session {
public:
DirectionSession(const SessionOptions& options, const Device* device_mgr, DirectSessionFactory* factory); Status Create(const GraphDef& graph) override;
Status Extend(const GraphDef& graph) override;
Status Run(...) override;//运行图 Status PRunSetup(...);//部分运行图准备
Status PRun(...);//部分运行图 Status Reset(const std::vector<string>& containers);//清空device_mgr中的containers,如果containers本身就是空的,那么清空默认容器 Status ListDevice(...) override;
Status Close() overrides;
Status LocalDeviceManager(const DeviceMgr** output) overrides; void ExportCostModels(...); private:
Status MaybeInitializeExecutionState(...);//给定graph之后,如果执行器状态没有初始化,则初始化基础的执行器状态 Status GetOrCreateExecutors(...);//对于一组给定的输入和输出,在一个给定的执行器集合中检索,是否存在合适的执行器,如果没有,则创造一个 Status CreateGraphs(...);//给定graph_def_和设备,以及输入和输出,创造多张图,这些新创建的图共享一个公共的函数库flib_def Status ExtendLocked(const GraphDef& graph);//Extend的内部执行类 Status ResourceHandleToInputTensor(...); Status SendPRunInputs(...);//将更多的输入提供给执行器,启动后续的执行 Status RecvPRunOutputs(...);//从执行器中获取更多的输出,它会等待直到输出张量计算完成 Status CheckFetch(...);//检查需求的输出能否根据给定的输入计算出来 Status WaitForNotification(...);
Status CheckNotClosed(); const SessionOptions options_; //设备相关的结构
const std::unique_ptr<const DeviceMgr> device_mgr_;
std::vector<Device*> devices_;
DeviceSet device_set_; string session_handle_;
bool graph_created_ GUARDED_BY(graph_def_lock_) = false;
mutex graph_def_lock_;
GraphDef graph_def_ GUARDED_BY(graph_def_lock_); std::vector<std::pair<thread::ThreadPool*, bool>> thread_pools_;//被用来执行op的线程池,用一个布尔值来标志,是否拥有这个线程池 Status init_error_; bool sync_on_finish_ = true;//如果为真,阻塞线程直到设备已经完成了某个步骤内的所有队列中的操作
void SchedClosure(thread::ThreadPool* pool, std::function<void()> c);//在线程池中调度c mutex executor_lock_;//保护执行器 std::unordered_map<string, std::shared_ptr<ExecutorsAndkeys>> executor_ GUARDED_BY(executor_lock_);//由签名映射到它的执行器,签名包括了部分执行图的输入和输出,由这两个就能唯一确定一个部分执行图 std::unordered_map<string, std::shared_ptr<RunState>> partial_runs_ GUARDED_BY(executor_lock_);//从签名到部分执行状态,每一个部分执行都会有一个专门保存其状态的结构 SessionState session_state_;//保存了所有当前在会话中正在存活的张量 DirectSessionFactory* const factory_;
CancellationManager* cancellation_manager_; std::unordered_map<string, string> stateful_placements_ GUARDED_BY(graph_def_lock_);//对于有状态的节点(比如params和queue),保存节点名称到节点所在设备的映射,一旦这些节点被放置在了某个设备上,是不允许再移动的 std::unique_ptr<SimpleGraphExecutionState> execution_state_ GUARDED_BY(graph_def_lock_);//放置整张图时使用 std::unique_ptr<FunctionLibraryDefinition> flib_def_;//在任何的重写或优化之前的函数库,特别是,CreateGraphs函数会修改函数库 mutex closed_lock_;
bool closed_ GUARDED_BY(closed_lock_) = false;//如果会话已经被关闭,则为true //为这个会话生成唯一的名字
std::atomic<int64> edge_name_counter_ = {0};
std::atomic<int64> handle_name_counter_ = {0}; static std::atomic_int_fast64_t step_id_counter_;//为所有的会话生成唯一的step id const int64 operation_timeout_in_ms_ = 0;//全局对阻塞操作的超时阈值 CostModelManager cost_model_manager_;//为当前会话中执行的图管理所有的损失模型
}

可见,DirectSession里面的很多内容都是为部分执行准备的。由于计算图仅是一个计算的规划,我们可以通过为同一张图选取不同的输入和输出,来执行不同的计算。而不同的计算需要不同的执行器,也需要不同的存储结构来保存各个计算的当前状态。为此,TF专门给出了几个结构体,首先我们来看一下对不同计算执行器的封装:

//为每一个partition准备的执行器和函数运行时库
struct PerPartionExecutorAndLib {
Graph* graph = nullptr;
std::unique_ptr<FunctionLibraryRuntime> flib;
std::unique_ptr<Executor> executor;
}; //为每一次计算提供的数据结构
struct ExecutorsAndKeys {
std::atomic_int_fast64_t step_count;
std::unique_ptr<Graph> graph;
NameNodeMap name_to_node;
std::unique_ptr<FunctionLibraryDefinition> flib_def;
std::vector<PerPartitionExecutorsAndLib> items;
std::unordered_map<string, size_t> input_name_to_index;
std::unordered_map<string, string> input_name_to_rendezvous_key;
std::unordered_map<string, size_t> output_name_to_index;
std::unordered_map<string, string> output_name_to_rendezvous_key; DataTypeVector input_types;
DataTypeVector output_types;
};

对于一张计算图来说,我们的每一次计算的执行,不论是完整图的计算还是部分图的计算,都有可能是跨设备的,因此都需要先做节点放置,把图的节点分割到不同的设备上,每一个设备上放置了一个图的partition,每个partition有对应的运行时函数库和执行器。而对于每一种计算来说,我们需要一个vector把不同partition的信息存储起来。

另外,刚才提到我们还需要为每一次计算提供保存当前状态的结构,下面就来看一下:

//对于每一个partition内的执行,会话保存了一个RunState
struct RunState {
mutex mu_;
Status status GUARDED_BY(mu_);
IntraProcessRendezvous* rendez = nullptr;
std::unique_ptr<StepStatsCollector> collector;
Notification executors_done;
std::unordered_map<string, bool> pending_inputs;//如果已经提供了输入,则为true
std::unordered_map<string, bool> pending_outputs;//如果已经获得了输出,则为true
TensorStore tensor_store;
ScopedStepContainer step-container;
//...
}; struct RunStateArgs {
RunStateArgs(const DebugOption& options) : debug_options(options) {}
bool is_partial_run = false;
string handle;
std::unique_ptr<Graph> graph;
const DebugOptions& debug_options;
};

其中,RunState为每一个partition的执行提供了状态保存的功能,而RunStateArgs则为前者提供了用于调试的参数和配置。

2.2 direct_session.cc

在源文件里,给出了DirectSessionFactory的定义,它提供了对于DirectSession进行生成和管理的功能,简要摘录如下:

class DirectSessionFactory : public SessionFactory {
public:
Session* NewSession(const SessionOptions& options) override;
Status Reset(...) override;
void Deregister(const DirectSession* session);
private:
mutex session_lock_;
std::vector<DirectSession*> session_ GUARDED_BY(sessions_lock_);//用于存储生成的DirectSession
};

另外,还提供了一个对于直接工厂注册的类:

class DirectSessionRegistrar {
public:
DirectSessionRegistrar() {
SessionFactory::Register("DIRECT_SESSION", new DirectSessionFactory());
}
};
static DirectSessionRegistrar registrar;

下面,我们会按照顺序对DirectSession内重要的函数,进行拆解,由于部分函数细节比较多,除了核心代码之外,我们仅给出功能解释:

DirectSession::DirectSession(const SessionOptions& options, const DeviceMgr* device_mgr, DirectSessionFactory* const factory){
//根据options准备线程池
//根据device_mgr准备device_和device_set_和每个设备的op_segment()
} Status DirectSession::Run(...){
//提取对于当前会话的本次运行的输入的名称
//检查对于所需的输入输出,是否已经存在现成的执行器
//构造一个调用帧(call frame),方便会话与执行器之间传递输入和输出
//创建一个运行时状态的结构(RunState)
//开始并行执行,核心代码如下
for(const auto& item : executors_and_keys->items){
item.executor->RunAsync(args, barrier->Get());
}
//获取输出
//保存本次运行中我们希望保存的输出张量
//创建并返回损失模型(cost model)
//如果RunOptions中有相关配置,输出分割后的图
} Status DirectSession::GetOrCreateExecutors(...){
//快速查找路径
//慢查找路径,对输入和输出做排序,使得相同输入和输出集合会得到相同的签名
//如果未找到,则创建这个执行器并缓存
//构建执行图,核心代码如下
CreateGraphs(options, &graphs, &ek->flib_def, run_state_args, &ek->input_types, &ek->output_types));
//为各子图准备运行时环境
} Status DirectSession::CreateGraphs(...){
//前期预处理
//图分割算法,核心代码如下
Partition(popts, &client_graph->graph, &partitions);
//检查分割结果的有效性
//图优化遍历,核心代码如下
OptimizationPassRegistry::Global()->RunGrouping(OptimizationPassRegistry::POST_PARTITIONING, optimization_options);
//允许设备重写它拥有的子图
}

可见,具体的执行过程是在Run函数内部,调用executor->RunAsync函数来实现的,在具体执行之前,我们还需要通过GetOrCreateExecutors函数获得执行器,在这个函数内部,我们通过CreateGraphs函数对原图进行了分割,并利用图优化遍历算法对图进行了优化。

tensorflow源码解析之common_runtime-direct_session的更多相关文章

  1. tensorflow源码解析之common_runtime拾遗

    把common_runtime中剩余的内容,按照文件名排序进行了简单的解析,时间原因写的很仓促,算是占个坑,后续有了新的理解再来补充. allocator_retry 有时候内存分配不可能一次完成,为 ...

  2. tensorflow源码解析系列文章索引

    文章索引 framework解析 resource allocator tensor op node kernel graph device function shape_inference 拾遗 c ...

  3. Tensorflow源码解析1 -- 内核架构和源码结构

    1 主流深度学习框架对比 当今的软件开发基本都是分层化和模块化的,应用层开发会基于框架层.比如开发Linux Driver会基于Linux kernel,开发Android app会基于Android ...

  4. tensorflow源码解析之framework拾遗

    把framework中剩余的内容,按照文件名进行了简单解析.时间原因写的很仓促,算是占个坑,后面有了新的理解再来补充. allocation_description.proto 一个对单次内存分配结果 ...

  5. tensorflow源码解析之common_runtime-executor-上

    目录 核心概念 executor.h Executor NewLocalExecutor ExecutorBarrier executor.cc structs GraphView ExecutorI ...

  6. tensorflow源码解析之framework-allocator

    目录 什么是allocator 内存分配器的管理 内存分配追踪 其它结构 关系图 涉及的文件 迭代记录 1. 什么是allocator Allocator是所有内存分配器的基类,它定义了内存分配器需要 ...

  7. tensorflow源码解析之common_runtime-executor-下

    目录 核心概念 executor.h Executor NewLocalExecutor ExecutorBarrier executor.cc structs GraphView ExecutorI ...

  8. tensorflow源码解析之distributed_runtime

    本篇主要介绍TF的分布式运行时的基本概念.为了对TF的分布式运行机制有一个大致的了解,我们先结合/tensorflow/core/protobuf中的文件给出对TF分布式集群的初步理解,然后介绍/te ...

  9. tensorflow源码解析之framework-op

    目录 什么是op op_def定义 op注册 op构建与注册辅助结构 op重写 关系图 涉及的文件 迭代记录 1. 什么是op op和kernel是TF框架中最重要的两个概念,如果一定要做一个类比的话 ...

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

    Session概述 1. Session是TensorFlow前后端连接的桥梁.用户利用session使得client能够与master的执行引擎建立连接,并通过session.run()来触发一次计 ...

随机推荐

  1. JVM学习六:堆分析

    一.内存溢出(OOM)的原因 在JVM中,有哪些内存区间? 堆溢出 public static void main(String args[]){ ArrayList<byte[]> li ...

  2. Linux发行版·常见Linux系统下载

    本专题页汇总最受欢迎的Linux发行版基本介绍和下载地址,如果您是一位刚接触Linux的新手,这里的介绍可能对您有所帮助,如果您是以为Linux使用前辈,也可以在评论处留下您宝贵意见和经验,以便让更多 ...

  3. socket 套接字编程

    今日内容 socket 套接字编程 简易服务端与客户端代码实现 通信循环 黏包现象(TCP协议) 报头制作.struct 模块.封装形式 内容详细 一.socket 套接字编程 实现一款能够进行数据交 ...

  4. 正则表达式以及re模块的使用

    内容概要 正则表达式简介 字符组 特殊符号 量词 贪婪匹配与非贪婪匹配 取消转义 正则表达式简介 '''正则表达式是一门语言,如果想在python中使用,需要导入re模块''' # 什么是正则表达式? ...

  5. kubeasz 部署高可用 kubernetes 集群

    文章目录 环境准备 配置模板机 配置hosts解析 配置ssh 免密钥登陆 kubeasz 部署服务准备 配置主机清单 部署集群 环境准备 IP HOSTNAME SYSTEM 192.168.131 ...

  6. Spring Boot数据访问之Druid连接池的配置

    在Spring Boot数据访问之数据源自动配置 - 池塘里洗澡的鸭子 - 博客园 (cnblogs.com)中数据源连接池方式缺省(type)情况下默认使用HikariCP,那不缺省如何配置呢?我们 ...

  7. 私有化轻量级持续集成部署方案--07-私有NPM仓库-Verdaccio

    提示:本系列笔记全部存在于 Github, 可以直接在 Github 查看全部笔记 对于个人来说,私有NPM仓库 作用性基本很小,但是对于企业,私有NPM仓库 可以保护代码暴露,具有很大的意义. 也是 ...

  8. BugKu-baby_flag.txt

    下载zip压缩包解压得到图片(你们不要再打啦!hhh)使用winhex或者010查看,最下面看到一串数字 怀疑可能是十六进制转换字符串,然后使用在线网站进行转换,得到一串base64编码, 再进行ba ...

  9. Java并发杂谈(一):volatile的底层原理,从字节码到CPU

    volatile的特性 volatile是Java中用于修饰变量的关键字,其主要是保证了该变量的可见性以及顺序性,但是没有保证原子性:其是Java中最为轻量级的同步关键字: 接下来我将会一步步来分析v ...

  10. RFC2544学习频率“Learning Frequency”详解—信而泰网络测试仪实操

    在RFC2544中, 会有一个Learning Frequency的字段让我们选择, 其值有4个, 分别是learn once, learn Every Trial, Learn Every Fram ...