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