『深度长文』Tensorflow代码解析(五)

深度学习大讲堂致力于推送人工智能,深度学习方面的最新技术,产品以及活动。请关注我们的知乎专栏!




6. TF – Session分析

客户端使用会话来和TF系统交互,一般的模式是,建立会话,在会话中添加节点和边,形成一张图,然后执行。如图 6 1所示,包含一个client和一个master。Client通过Session与Master交互,而Master协调一个或者多个的 worker processes 完成计算任务,每个 worker process 负责对一个或者多个计算设备(CPU 核或者 GPU card)的任意访问和在这些设备上进行图节点的计算按照 master 的要求执行。

图 6 1中Session.run()向Master发送计算请求,执行graph计算。Session.run函数定义如下:
TF中Session有两种实现:一种是DirectSession类,可以执行localhost环境运行的图;另一种是GrpcSession类,可以执行分布式环境运行的图。

6.1 Direct Session分析

DirectSession定义在localhost环境,运行环境如下所示。图 6 2中仅一个本地worker,由master管理worker完成计算任务。

首先,从run函数开始,指定feed节点、fetch节点和feed数据。
创建Executors。Executor将执行graph计算操作,多个Executor可以并行计算,在feed批处理计算数据时非常有用。
将input放入Rendez消息信箱中,等待consumer取出。
创建ExecutorBarrier。ExecutorBarrier在多个Executor并行计算时起协调作用,保证每个Executor执行graph计算时数据的一致性。
在barrier的协调下,每个executor完成对应的graph计算操作。
每个executor对应一个ExecutorState实例。executorstate跟踪graph中的node,通常认为node有完成计算、准备计算、等待计算等状态。当node处于就绪状态就为该节点执行计算操作。
创建节点序列触发器TaggedNodeSeq 变量,即ready队列。先为ready队列添加root节点,即起始点,然后放入schedule任务队列中。
从ready队首开始,执行计算操作,计算过程由ExecutorState::Process函数完成。
ExecutorState::Process调用每个节点中OpKernel对应的Compute函数或ComputeAsync函数,如MatMulOp实现了Compute函数,RecvOp实现了ComputeAsync函数。
6.2 Grpc Session分析

Grpc Session是TF分布式环境的会话模式,run函数向master发送请求后,master会启动两个或两个以上的worker完成计算任务。本章节先简单介绍TF的分布式架构,然后从代码角度逐步分析。

6.2.1 分布式架构分析

分布式架构涉及到任务调度、数据传输、节点进程监控、容错机制等。TF分布式环境采用master-slave架构,由master service完成主控操作。

Ø 任务调度。分布式系统会根据节点设备负载状态分配任务,达到系统负载均衡的状态。TF目前没有任务调度功能,但通过设备分配算法部分解决了设备负载问题。

Ø 数据传输。TF实现了Rendezvous消息机制,如图 6 3,由发送节点(SendOp)和接收节点(RecvOp)为不同设备的两个相邻节点完成完成数据通信操作。Tensor数据封装成protobuf格式,数据传输由gRPC完成。

Ø 节点进程监控。分布式系统通常使用周期性给slave节点发送心跳文件监控节点状态。TF中master进程采用类似的方法对worker节点进行周期性检查。

Ø 容错机制。分布式系统的容错机制允许数据传输断开、节点进程异常等。当错误产生时,整个运算流图的运行将会终止并再重头开始。TensorFlow通过流图运行过程中与张量关联的Variable节点,实现重启中状态和核对点的稳定恢复。每个Variable节点与Save节点相关联。Save节点周期运行(每N此迭代运行一次,或每N秒运行一次)。当Save节点运行时,Variable变量的内容被写入非易失性存储器中(如分布式文件系统中)。同样地,每个Variable变量与Restore节点关联,Restore节点仅在重启后的第一次迭代运行中生效。

TF分布式环境包含多个worker节点,分布在不同的设备中。每个worker节点是一个graph子图,利用SendOp/RecvO完成跨worker进程间通信来实现数据跨设备传输。

TF Graph设计高度灵活,使分布式训练框架设计也灵活多样。TF分布式模型有多设备模型串行、同步模型并行、异步模型并行、混合模型并行。

模型串行只有一个graph副本,可以将不同的layer分配当不同的设备上,也可以将同一layer分配到不同设备上。举例来说,可以将两层RNN网络部署在两块GPU卡或两个GPU服务器上,也可以将单层RNN网络划分给两个设备执行。这种模型设计本质上类似于单机模型。

模型并行会产生多个graph副本,分布在多个worker节点,还需要一个或多个ps节点负责参数更新,如图 6 4所示[1]。

同步模型并行训练(图 6 4上子图)中,每个worker节点先从ps节点获取参数数据更新到worker本地副本,再接受一个batch输入数据,神经网络模型通过正向计算得到损失函数的损失值,再通过反向计算得到所有参数的梯度更新值,将梯度更新值发送给ps节点。ps节点等待并接收所有worker节点发送的梯度更新值,对梯度更新值作求和或均值处理后利用optimizer更新所有参数值,再同步到每个worker节点的参数副本,如此循环训练。

异步模型并行训练(图 6 4下子图)中,每个client通过相应的session与master交互。每个worker节点从ps节点获取参数数据更新到worker本地副本,worker的处理过程与同步训练是一致的,不同的是ps节点处理方式。ps节点在得到某个worker的参数梯度更新值后,不会等待其他worker发送,直接利用梯度更新值更新梯度,并将更新的参数同步到对应的work的节点。可以看出,每个worker与ps构成独立完整的闭环,每个worker耦合度很低,worker之间参数梯度关联性也不强,训练速度虽然比同步模型快,但收敛速度却不及同步模型。

6.2.2 分布式代码分析

TF分布式会话交互由类GrpcSession发起,类GrpcMasterService调度,GrpcWorkerService类执行计算任务,这三个类分别位于distributed_runtime/rpc/grpc_session.cc、distributed_ runtime/rpc/grpc_master_service.cc、distributed_runtime/rpc/grpc_worker_service.cc文件中,如图 6 5所示[17]。后面distributed_runtime均简写为dist。

下面讲述Grpc Session的执行过程,图 6 6中是Grpc的函数调用结构图,可以大略看出从session到master,再到worker的计算流程。
首先,从GrpcSession::Run函数开始,即声明feed和fetch。这里变量feed是std::vector<std::pair<string, Tensor>>类型,fetch是std::vector<string>类型,output是std::vector<Tensor>类型。
在RunStepRequest中加入feed和fetch,然后通过protobuf发送给remote master。
GrpcRemoteMaster接收到grpc session的请求,转交给grpc master service,经历了GrpcSession -> GrpcRemoteMaster -> GrpcMasterService -> Master -> MasterSession的任务转交。
MasterSession执行计算任务分为两步:RegisterPartitions将 graph按location不同切分为多个subgraph;RunPartitions执行subgraph。这里的一个subgraph就是一个worker,对应一个worker service。
RunPartitions函数执行计算操作时,为每个subgraph配置一个worker。
每个worker节点调用RunGraphAsync执行subgraph的计算任务。这里subgraph是通过protobuf协议给worker service发送计算请求的。
GrpcRemoteWorker接收master的请求,交由GrpcWorkerService执行。
Worker service把计算任务放进线程池队列中,计算任务在DoRunGraph函数中。
GrpcWorkerService::DoRunGraph函数依赖GraphMgr::ExecuteAsync完成worker的计算。ExecuteAsync计算过程与6.1节中DirectSession::Run()执行过程类似,即创建Executors执行计算操作,创建ExecutorBarrier协调各个Executor,创建ExecutorState跟踪node状态。
GrpcSession相关测试文件在distributed_runtime/master_test.cc,调试方法:
7. Reference

[1]. Abadi M, Agarwal A, Barham P, et al. Tensorflow: Large-scale machine learning on heterogeneous distributed systems[J]. arXiv preprint arXiv:1603.04467, 2016.

[2]. Tensorflow in Github. tensorflow/tensorflow.

[3]. Tensorflow official website. tensorflow.org/

[4]. TensorFlow™ - Open Source Library for Machine Learning Applications. https://delftswa. gitbooks.io/desosa2016/content/tensorflow/chapter.html

[5]. Eigen unsupported documentation. Eigen's unsupported modules index.html

[6]. Tensor official documentation. eigen / eigen / source / unsupported Eigen/CXX11/src/Tensor/README.md?fileviewer=file-view-default

[7]. 如何评价Tensorflow和其它深度学习系统. Sina Visitor System 75666.

其中评估时间的函数EstimateComputationCosts是对graph中每个node依次评估,节点计算时间评估函数如下。

[8]. Eigen official documentation. Eigen: Main Page

[9]. Eigen Benchmark. Benchmark - Eigen

[10]. Half-precision floating-point format. en.wikipedia.org/wiki/H -point_format

[11]. Tensorflow设备内存分配算法解析. Sina Visitor System

[12]. Tensorflow: Adding a New Op. tensorflow.org/how_tos/ -a-new-op

[13]. Tensorflow: Reading data. tensorflow.org/how_tos/

[14]. Tensorflow: Custom Data Readers. tensorflow.org/how_tos/

[15]. Goldsborough P. A Tour of TensorFlow[J]. 2016.

[16]. Click C. Global code motion/global value numbering[C]//ACM SIGPLAN Notices. ACM, 1995, 30(6): 246-257.

[17]. Kevin Robinson. A tour through the TensorFlow codebase.


该文章属于“深度学习大讲堂”原创,如需要转载,请联系@果果是枚开心果.

作者简介:

姚健,毕业于中科院计算所网络数据实验室,毕业后就职于360天眼实验室,主要从事深度学习和增强学习相关研究工作。目前就职于腾讯MIG事业部,从事神经机器翻译工作。联系方式: yao_62995@163.com

原文链接:『深度长文』Tensorflow代码解析(五)

欢迎大家关注我们的微信公众号,搜索微信名称:深度学习大讲堂

发布于 2017-03-27

文章被以下专栏收录

    推送深度学习的最新消息,包括最新技术进展,使用以及活动,由中科视拓(SeetaTech)运营。 中科视拓目前正在招聘: 人脸识别算法研究员,深度学习算法工程师,GPU研发工程师, C++研发工程师,Python研发工程师,嵌入式视觉研发工程师,PR经理,商务经理。(PS:深度学习算法工程师岗位、Python研发工程师岗位、嵌入式视觉开发工程师岗位和运营岗位同时接收实习生投递)有兴趣可以发邮件至:hr@seetatech.com,想了解更多可以访问,www.seetatech.com