GPU 设备插件

TensorFlow 的 可插拔设备 架构将新的设备支持添加为单独的插件包，这些插件包与官方 TensorFlow 包一起安装。

此机制不需要对 TensorFlow 代码进行任何特定于设备的更改。它依赖于 C API 以稳定方式与 TensorFlow 二进制文件通信。插件开发人员为其插件维护单独的代码存储库和分发包，并负责测试其设备。

使用设备插件

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

# Install the APU example plug-in package
$ pip install tensorflow-apu-0.0.1-cp36-cp36m-linux_x86_64.whl
...
Successfully installed tensorflow-apu-0.0.1

安装插件后，测试设备是否可见，并在新 APU 设备上运行操作

import tensorflow as tf   # TensorFlow registers PluggableDevices here.
tf.config.list_physical_devices()  # APU device is visible to TensorFlow.
[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.

with tf.device("/CPU:0"):
  c = tf.nn.relu(a)        # Runs on CPU.

@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.

可用设备

适用于 macOS GPU 的 Metal PluggableDevice

适用于 Windows 和 WSL 的 DirectML PluggableDevice（预览版）

适用于 Linux 和 WSL 的 Intel® Extension for TensorFlow PluggableDevice