Tenbsorflow设备注册系列——（一）Device概要梳理

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 数组中