作者：Penporn Koanantakool 和 Pankaj Kanwar。

随着机器学习生态系统中加速器（GPU、TPU）数量的激增，人们迫切需要将新的加速器与 TensorFlow 无缝集成。在这篇文章中，我们介绍了 可插拔设备 架构，该架构提供了一种插件机制，用于将设备注册到 TensorFlow，而无需更改 TensorFlow 代码。

此可插拔设备架构是在 TensorFlow 社区内协作设计和开发的。它利用了为 模块化 TensorFlow 完成的工作，并使用 StreamExecutor C API 构建。可插拔设备机制在 TF 2.5 中可用。

对无缝集成的需求

在此之前，任何新设备的集成都需要对 TensorFlow 核心进行更改。由于以下几个问题，这种方法不可扩展，例如

• 复杂的构建依赖项和编译器工具链。 引入新的编译器并非易事，并且会增加产品的技术复杂性。
• 开发时间缓慢。 更改需要经过 TensorFlow 团队的代码审查，这可能需要时间。增加的技术复杂性也会增加新功能的开发和测试时间。
• 要测试的构建配置数量组合。 为特定设备进行的更改可能会影响其他设备或 TensorFlow 的其他组件。每个新设备都可能以乘法方式增加测试配置的数量。
• 易于破坏。 缺乏通过定义明确的 API 制定的契约意味着更容易破坏特定设备。

什么是可插拔设备？

可插拔设备机制不需要对 TensorFlow 代码进行特定于设备的更改。它依赖于 C API 以稳定方式与 TensorFlow 二进制文件进行通信。插件开发者维护他们插件的单独代码库和分发包，并负责测试他们的设备。这样，TensorFlow 的构建依赖项、工具链和测试过程就不会受到影响。集成也更不容易出错，因为只有对 C API 或可插拔设备组件的更改才会影响代码。

可插拔设备机制有四个主要组件

• 可插拔设备类型： TensorFlow 中的 新设备类型，允许从插件包中注册设备。在设备放置阶段，它优先于原生设备。
• 自定义操作和内核： 插件通过 内核和操作注册 C API 向 TensorFlow 注册他们自己的操作和内核。
• 设备执行和内存管理： TensorFlow 通过 StreamExecutor C API 管理插件设备。
• 自定义图优化过程： 插件可以通过 图优化 C API 注册一个自定义图优化过程，该过程将在所有标准 Grappler 过程之后运行。

设备插件如何与 TensorFlow 交互。

使用可插拔设备

为了能够使用特定设备，就像在 TensorFlow 中使用原生设备一样，用户只需安装该设备的设备插件包即可。以下代码片段显示了如何安装和使用新设备（例如 Awesome Processing Unit (APU)）的插件。为简单起见，假设此 APU 插件只有一个用于 ReLU 的自定义内核。

$ pip install tensorflow-apu-0.0.1-cp36-cp36m-linux_x86_64.whl
…
Successfully installed tensorflow-apu-0.0.1
$ python
Python 3.6.9 (default, Oct  8 2020, 12:12:24) 
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf  # TensorFlow registers PluggableDevices here
>>> tf.config.list_physical_devices()
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:APU:0', device_type='APU')]

>>> a = tf.random.normal(shape=[5], dtype=tf.float32)  # Runs on CPU
>>> b =  tf.nn.relu(a)  # Runs on APU
 
>>> with tf.device("/APU:0"):  # Users can also use 'with tf.device' syntax
...  c = tf.nn.relu(a)  # Runs on APU
 
>>> @tf.function  # Defining a tf.function
... def run():
...    d = tf.random.uniform(shape=[100], dtype=tf.float32)  # Runs on CPU
...    e = tf.nn.relu(d)  # Runs on APU
>>> run()  # PluggableDevices also work with tf.function and graph mode.

即将推出的可插拔设备

我们很高兴宣布，英特尔将成为我们的第一个发布可插拔设备的合作伙伴。英特尔为这项工作做出了重大贡献，提交了超过 3 个 RFC 来实现整个机制。他们将发布一个 英特尔扩展 TensorFlow 插件包，将英特尔 XPU 带到 TensorFlow 中，以加速 AI 工作负载。我们也期待其他合作伙伴利用可插拔设备并发布其他插件。

我们将发布详细的教程，说明如何为可能对利用此基础设施感兴趣的合作伙伴开发可插拔设备插件。有关可插拔设备的问题，工程师可以直接在 RFC PR 上发布问题 [1, 2, 3, 4, 5, 6], 或者在 TensorFlow 论坛 上使用标签 pluggable_device 发布问题。