GPUデバイスプラグイン

TensorFlow のプラガブル デバイスアーキテクチャは、公式 TensorFlow パッケージと一緒にインストールされる別個のプラグイン パッケージとして新しいデバイス サポートを追加します。

このメカニズムでは、TensorFlow コードにデバイス固有の変更を加える必要はありません。 C API に依存して、安定した方法で TensorFlow バイナリと通信します。プラグイン開発者は、プラグイン用に個別のコード リポジトリと配布パッケージを管理し、デバイスのテストを担当します。

デバイスプラグインを使用する

TensorFlow のネイティブ デバイスと同様に、特定のデバイスを使用するには、ユーザーはそのデバイスのデバイス プラグイン パッケージをインストールするだけで済みます。次のコード スニペットは、新しいデモンストレーション デバイスのプラグインであるAwesome Processing Unit (APU)がどのようにインストールされ、使用されるかを示しています。わかりやすくするために、このサンプル APU プラグインには ReLU 用のカスタム カーネルが 1 つだけ含まれています。

# 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 用 TensorFlow PluggableDevice用インテル® 拡張機能: