1、概述
- TensorFlow 中的 Device 是通过注册机制添加到运行的进程中的
- TensorFlow 使用工厂来创建各种各样的 Device,并且几乎为每一种 Device 都实现了对应的 DeviceFactory
- 在定义每个Device时,通过利用C++事先定义好的宏(Macro)将类对象主动注册到工厂中
- 这样就可以达到在程序启动完毕时,工厂里已经储备有各种各样所需要的内容
2、Device 类图梳理
- DeviceAttributes:
- 在 TensorFlow 源码中并不能直接找到这个类的 c++ 定义
- 它是由protobuf编译出来的
- 其含义是对特定Device属性的封装,比如 Device 的 type,存储的限制等等
- DeviceBase:
- 定义了Device用到的基本方法
- 比较重要的是 GetAllocator 和 MakeTensorFromProto
- 前者返回存储器的分配器
- 后者是从 Proto 中生成 Tensor,该方法必须被重写;
- Device:
- 这个类比 DeviceBase 更加具体,新包含了一些用于计算调用的方法
- 比如 Compute 函数就会调用某 Op 的 Compute 计算
- RemoteDevice:
- 这个类会在分布式时使用
- 具体功能待补充、
- SingleThreadedCpuDevice:
- 这是一个仅有单个线程的 CPU Device 抽象
- 它和 ThreadPoolDevice 不同,只被用于 in expensive 的 Op 计算
- 这样做的好处是避免了一些 thread 初始化工作
- ThreadPoolDevice:
- CPU Device 的实现
- RenamedDevice:
- Device的封装类
- 封装时会再取一个新的 Device name
- GPUCompatibleCPUDevice:
- 这也是 CPU Device 的实现
- 和 ThreadPoolDevice 不同的是,它更多的是为了和 GPU 进行交互而存在
- 从其使用的 CudaHostAllocator 就可以看出这一点
- BaseGPUDevice
- 和 GPU Device 的实现有关
- GPUDevice:
- 在继承 BaseGPUDevice 的基础上重写了 Allocator
3、DeviceFactory 类图梳理
- 概述
- TensorFlow在启动时会调用一系列static的函数
- 这些函数是通过宏(Macro)展开得到的
- 对于设备管理模块来说,每种 Device 都由对应的 Factory 负责管理
- 而每种 DeviceFactory 会在程序启动时注册到全局唯一的 static device factory 表中
- Registrar
- 可以看做是一个带有模板的控制类
- 它只负责一件事——各种DeviceFactory向全局表的注册
4、DeviceFactory 注册机制
- 流程
- 在代码层面,注册函数的调用是通过宏实现的
- 该宏通过传入 DeviceFactory 的类名称即可触发Regsitrar的调用逻辑
- 在每个 DeviceFactory 的 C++ 实现文件后面都会引用此宏
- 有了上述的 DeviceFactory 的注册后,就可以在使用时根据使用的Device类型,从对应的 DeviceFactory中“获取”想要的 Device 了
5、Device 创建
- 流程
- 真正从 Factory 中取出 Device 的过程是在 Session 创建时进行的
- 调用的函数是 DeviceFactory 中的 static 函数 AddDevices
- AddDevices
- 它会遍历全局 device factories 表中全部的 DeviceFactory 并取出
- 然后逐个调用每个具体 XXDeviceFactory 的 CreateDevices 函数
- 将创建的 Device 放进 vector 数组中
- PS
- 实际上 DeviceFactory 的 static 函数调用 AddDevices 时,会先将CPU Device创建出来
- 如果没有可用的 CPU Device,那么程序就会直接报错退出
- 这是因为 TensorFlow 需要保证当没有其他 Device 存在时,至少还有CPU可以完成整体程序的计算和调度运行
- 创建 CPU Device 之后,就会去遍历所有 DeviceFactory
- 然后把所有能够创建的 Device 全部创建出来放入 vector 数组中