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

背景

作为一款优秀的异构深度学习算法框架,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函数CreateDevices。它会遍历全局device factories表中全部的DeviceFactory并取出,然后逐个调用每个具体XXDeviceFactory的CreateDevices函数,将创建的Device放进vector数组中。下面给出一个简化版的时序图。

上述的时序图描述的较为简单,实际上DeviceFactory的static函数调用CreateDevices时会先将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中补充。

猜你喜欢

转载自www.cnblogs.com/deep-learning-stacks/p/9313700.html