TensorFlow中的设备管理——Device的创建与注册机制

背景

[作者：DeepLearningStack，阿里巴巴算法工程师，开源TensorFlow Contributor]

作为一款优秀的异构深度学习算法框架，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时才会触发，该过程十分简单。