tensorflow源码分析（六）设备管理

2018年12月20日 09:13:14 hobertony_7 阅读数：61

接下来介绍一下tensorflow对设备devices的发现和管理，负责运行kernel的具体硬件设备抽象成为devices；实现某种操作的算法叫做kernel，可以将kernel看做是一段能够跑在具体硬件设备上的算法程序，所以相同的算法实现，在gpu上有一套实现，在cpu上有一套实现。

在session初始化时会调用add_device接口，把cpu gpu设备添加到session到devices_表中。

上图主要介绍再本地运行时tensorflow内部对设备的抽象模型以及在源码中对设备抽象的类的继承关系，非本地的设备抽象类时remotedevice，这里不再详细介绍。

GPUCompatibleCPUDevice实例：

1)将 cpu 的计算资源抽象为thread pool，以支持多thread之间的并发执行；

2)将主机内存抽象为 CPUAllocator（分配器) 实例来进行管理，为cpu kernel、gpu kernel提供主机内存的申请、释放功能；

(如何对cpu资源抽象可以进一步研究的点)

GPUDevice实例：

1)将gpu的计算资源抽象为streams，由于目前只支持NVIDIA的gpu，所以这里我们可以看作抽象为cuda streams，多个cuda streams之间的计算可以并发处理；

2)通过GPUBFCAllocator实例来管理显存，为gpu kernel提供显存的申请、释放功能。

（gpu资源管理以及与cuda库的配合可以作为可以继续研究的点）

Tensorflow系统中可用的device实例在session初始化时创建，并由该session使用，归属于该session，device的创建使用的是factory模式。

该图主要描述了graph运行时时如何与device的抽象类交互，devices_是在session初始化时构建生成，里面包含了该session所有的可用设备，executor执行器会把graph中相应得node交给设备运行，其实就是在设备上执行node对应得opkernel。

==================================================================================================================================================================================================================================================================================================================================================================

Tensorflow中Device实例的生成和管理

2017年11月15日 19:55:42 jacob_wjj 阅读数：937

关键术语描述

kernel

在神经网络模型中，每个node都定义了自己需要完成的操作，比如要做卷积、矩阵相乘等。
可以将kernel看做是一段能够跑在具体硬件设备上的算法程序，所以即使同样的2D卷积算法，我们有基于gpu的Convolution 2D kernel实例、基于cpu的Convolution 2D kernel实例。

device

负责运行kernel的具体硬件设备抽象。每个device实例，对应系统中一个具体的处理器硬件，比如gpu:0 device, gpu:1 device, cpu:0 device。一般来说，每个device实例同时包括处理器资源、内存资源。device的抽象支持硬件设备提供的并行处理能力。

device是什么

为方便描述，下面我们把在tensorflow里面运行的神经网络模型都统一称为graph。

我们知道，tensorflow主要针对的是跨硬件平台、分布式、并发运行的场景，参与运算的每个硬件资源，我们都抽象为device实例，便于管理。

device的主要职责：

管理处理器资源，为支持device内部的并行计算，进一步将其抽象为thread pool或streams：
cpu：使用thread pool来管理，thread之间可支持不同程度的并行计算能力
gpu: 针对nvidia gpu，使用cuda streams来管理，根据不同的gpu型号，可支持不同数量的stream做并行计算

管理内存资源：为kernel的运行，分配和释放内存，进一步抽象为Allocator及其各种子类的实例来管理。

主机内存：
cpu kernel 计算时需要的内存。
gpu kernel的输出结果如果要放置到主机内存中时，gpu kernel也需要申请主机内存。

显存： gpu kernel 计算时需要的内存。

device的种类及应用场景

由于device要抽象的设备种类较多，我们主要描述一下本地运行的cpu device、gpu device实例类型。先用一个UML图来表示一下各种device抽象类的关系：

可以看到，cpu device实例使用的类是GPUCompatibleCPUDevice，主要是在ThreadPoolDevice的基础上，增加了gpu<-> cpu之间内存传输数据的优化措施。

gpu device实例使用的类是 GPUDevice 。
4. device实例的关键数据结构

我们以常用的cpu device，gpu device为例，用下图描述一下device实例的关键数据结构：

可以看到每个device实例内部都具备并行处理的能力：

GPUCompatibleCPUDevice实例
    将 cpu 的计算资源抽象为thread pool，以支持多thread之间的并发执行；
    将主机内存抽象为 CPUAllocator 实例来进行管理，为cpu kernel、gpu kernel提供主机内存的申请、释放功能；

GPUDevice实例
    将gpu的计算资源抽象为streams， 由于目前只支持NVIDIA的gpu，所以这里我们可以看作抽象为cuda streams，多个cuda streams之间的计算可以并发处理；
    通过GPUBFCAllocator实例来管理显存，为gpu kernel提供显存的申请、释放功能。

device实例的创建

系统中可用的device实例，由session发起创建，归属于session实例。

device的创建，使用Factory 设计模式，session会调用所有注册的device factory，逐一产出符合条件的device实例。

以DirectSession实例创建gpu device、cpu device为例，具体流程如下图所示。
为方便结合代码阅读，已包含主要的类、函数调用路径：

可以看到，最终产出的gpu device、cpu device实例，都会保存至DirectSession实例的 devices_ 表中，由DirectSession实例进行分配和使用。
6. 在graph运行阶段device的使用

在graph的创建阶段，session为每个node分配一个具体的device实例，同时为每个node创建一个具体的kernel实例，这个kernel实例将会运行在分配的device实例上。(参见Tensorflow 核心流程剖析 2 – 神经网络模型的创建和分割)

接下来，在graph的运行阶段，session会依次处理graph中的node，调度node所分配的device实例，去运行node的kernel实例。

每个kernel 在运行时，会向其分配的device，申请需要的计算资源、内存资源等，完成具体的运算操作。

上述流程如下图所示。
为方便结合代码阅读，已包含主要的类、函数调用路径：

TensorFlow中的设备管理（一）——Device的创建与***制

2018年07月16日 09:26:51 qq_42606051 阅读数：265

背景

作为一款优秀的异构深度学习算法框架，TensorFlow可以在多种设备上运行算法程序，包括CPU，GPU，Google开发的TPU等。因为TensorFlow的架构特性非常好，可扩展性很强，所以也支持用户自定义补充其他计算设备，比如可以接入FPGA甚至是自定义芯片等。虽然在Google发布的TensorFlow white paper中并没有过多的描述设备管理相关的内容，只是从较高层面上阐述了Device以及Job的命名规则，但是其设备管理模块确实是处于架构中比较核心的地位。本文将从架构层面出发，详细阐述当前TensorFlow源码中关于设备管理的设计思想和相关细节，理解这部分内容不但可以加深对TensorFlow源码的理解，还可以有能力接入一些自定义的设备。本文是TensorFlow设备管理的第一篇文章，为了能让读者更好的切入到TensorFlow源码阅读过程中，先从较为简单地Device的创建和***制开始。读者也可以一边对照本文一边对照源码进行阅读和梳理，并欢迎大家提出各种相关的意见和建议。

计算设备（Device）定义

Google在2015年发布的第一版TensorFlow white paper中，从功能角度上阐述了Device的相关内容，我们可以总结出关键的几点如下：

1. 在TensorFlow中的Device有着特殊的命名规则，无论在单机还是分布式任务中，都能依靠命名确定唯一的Device，它是Device的唯一标识符；

2. TensorFlow使用***制将实现多种Device的添加管理；

3. 每个Device自己管理Memory的分配和释放。

Device的命名一般使用/job:{job_name}/task:{job_id}/device:{type}:{device_id}的格式，这是为了更好的支持分布式任务。例如/job:worker/task:17/device:gpu:3就表示该Device是ID为17的worker上的ID为GPU设备。至于分布式中的相关概念会在其他文章中详细阐述，在这里我们只需要知道Device的命名可以帮助我们在任务中定位到具体的某个唯一Device即可。

Device的***制

TensorFlow有两处涉及到了设备管理，一处存在于TensorFlow的core中，另一部分存在于XLA中。本文只会阐述TensorFlow core中的内容，关于XLA部分的讲解可以参见其他blog。TensorFlow使用工厂来创建各种各样的Device，并且几乎为每一种Device都实现了对应的DeviceFactory。初读代码时可能会被各种Device类名搞混，下面先从TensorFlow中已经有的Device类出发，给出各种Device的类说明。

Device相关类图

TensorFlow对不同种类的Device做了多层级的抽象，下面的类图是从当前TensorFlow源码中梳理出的比较重要的部分。

上图中的每个类（class）都可以在TensorFlow的源码中找到，因为当前TensorFlow的进化过程比较快，代码结构并不处于一个十分稳定的状态，所以上述类图中的类结构关系可能在未来发生一些变化，这一点从注释中也可以看出一些端倪，但是大的架构不会发生变化，所以梳理类结构也是十分有意义的。下面将对每个类的作用进行简单地阐述，读者在理解这些类的作用以及关系后再去阅读源码就会非常清晰了。

1. DeviceAttributes：在TensorFlow源码中并不能直接找到这个类的c++定义，其实它是由protobuf编译出来的。其含义也很好理解，是对特定Device属性的封装，比如Device的type，存储的限制等等；

2. DeviceBase：定义了Device用到的基本方法，比较重要的是GetAllocator和MakeTensorFromProto，前者返回存储器的分配器，后者是从Proto中生成Tensor，该方法必须被重写；

3. Device：这个类比DeviceBase更加具体，新包含了一些用于计算调用的方法，比如Compute函数就会调用某Op的Compute计算；

4. RemoteDevice：这个类会在分布式时使用，在此暂时不进行阐述；

5. SingleThreadedCpuDevice：这是一个仅有单个线程的CPU Device抽象，它和ThreadPoolDevice不同，只被用于in expensive的Op计算，这样做的好处是避免了一些thread初始化工作；

6. ThreadPoolDevice：这就是CPU Device的实现；

7. RenamedDevice：Device的封装类，封装时会再取一个新的Device name；

8. GPUCompatibleCPUDevice：这也是CPU Device的实现，和ThreadPoolDevice不同的是，它更多的是为了和GPU进行交互而存在，从其使用的CudaHostAllocator就可以看出这一点；

9. BaseGPUDevice，GPUDevice：这两个类都和GPU Device的实现有关，其中GPU Device类只是在继承BaseGPUDevice的基础上重写了Allocator，但没有理解这样设计的深层次原因。

DeviceFactory相关类图

上文提到过，TensorFlow中的Device是通过***制添加到运行的进程中的。***制在开源代码中是十分常见的设计技巧，它涉及到了一种非常经典的设计模式——工厂模式。在定义每个Device时，通过利用C++事先定义好的宏（Macro）将类对象主动注册到工厂中，这样就可以达到在程序启动完毕时，工厂里已经储备有各种各样所需要的内容。在TensorFlow中存在多处使用工厂模式的例子，比如本文阐述的Device管理，以及Session管理等。在其他开源框架中我们也能够看到这一模式，比如Caffe中的Layer也使用的是工厂模式。

从源码中可以看到，TensorFlow在启动时会调用一系列static的函数，这些函数是通过宏（Macro）展开得到的。对于设备管理模块来说，每种Device都由对应的Factory负责管理，而每种DeviceFactory会在程序启动时注册到全局唯一的static device factory表中。下面的类图展示了各种DeviceFactory之间的继承关系。

DeviceFactory的***制

在理清了上一小节中Device相关类的继承关系和说明之后，对于上图中各种DeviceFactory之间的继承关系就很好理解了。这里面比较陌生的类是Registrar，这可以看做是一个带有模板的控制类，它只负责一件事——各种DeviceFactory向全局表的注册。在代码层面，注册函数的调用是通过宏实现的，该宏通过传入DeviceFactory的类名称即可触发Regsitrar的调用逻辑，在每个DeviceFactory的C++实现文件后面都会引用此宏。下图形象的展示了注册的过程。

有了上述的DeviceFactory的注册后，就可以在使用时根据使用的Device类型，从对应的DeviceFactory中“获取”想要的Device了。

Device的创建

有了DeviceFactory之后，我们就可以从Factory中拿到各种各样的Device了。真正从Factory中取出Device的过程是在Session创建时进行的，调用的函数是DeviceFactory中的static函数AddDevices。它会遍历全局device factories表中全部的DeviceFactory并取出，然后逐个调用每个具体XXDeviceFactory的CreateDevices函数，将创建的Device放进vector数组中。下面给出一个简化版的时序图。

上述的时序图描述的较为简单，实际上DeviceFactory的static函数调用AddDevices时会先将CPU Device创建出来，如果没有可用的CPU Device，那么程序就会直接报错退出（一般情况下不会发生此类情况）。这是因为TensorFlow需要保证当没有其他Device存在时，至少还有CPU可以完成整体程序的计算和调度运行。创建CPU Device之后，就会去遍历所有DeviceFactory，把所有能够创建的Device全部创建出来放入vector数组中。

总结

本文主要阐述了TensorFlow设备管理模块中的设备创建于***制，它是TensorFlow进行设备管理的第一步，也是最简单的部分。想要深入TensorFlow源码的新人可以先从此模块开始阅读，进而熟悉TensorFlow的Coding style。Device的创建和注册过程触发于程序运行的初始阶段，因为创建Device时使用了工厂模式，所以此处涉及到了各种DeviceFactory的定义和注册。在TensorFlow的C++代码中，各种DeviceFactory在实现文件中通过宏主动将自己注册到全局表中，这样做的目的不但减少了大量重复的注册代码，还与Device的创建解耦合，是一个非常经典常见的编码技巧。至于Device的创建是在创建Session时才会触发，这个过程十分简单。至于设备管理模块中涉及到的其他内容将在后续blog中补充。

To help machines to understand the world.

设备