GPU 裝置外掛程式

TensorFlow 的可插拔裝置架構新增了裝置支援，以獨立外掛程式套件的形式與官方 TensorFlow 套件一同安裝。

此機制不需要在 TensorFlow 程式碼中進行裝置專屬變更。它仰賴 C API 以穩定的方式與 TensorFlow 二進位檔通訊。外掛程式開發人員維護其外掛程式的獨立程式碼存放區和發布套件，並負責測試其裝置。

使用裝置外掛程式

若要使用特定裝置 (如同在 TensorFlow 中使用原生裝置)，使用者只需安裝該裝置的裝置外掛程式套件即可。以下程式碼片段示範如何安裝及使用新示範裝置「超棒處理單元 (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