Plugin perangkat GPU

Arsitektur perangkat pluggable TensorFlow menambahkan dukungan perangkat baru sebagai paket plug-in terpisah yang diinstal bersama paket TensorFlow resmi.

Mekanisme ini tidak memerlukan perubahan khusus perangkat pada kode TensorFlow. Ini bergantung pada C API untuk berkomunikasi dengan biner TensorFlow secara stabil. Pengembang plugin memelihara repositori kode dan paket distribusi terpisah untuk plugin mereka dan bertanggung jawab untuk menguji perangkat mereka.

Gunakan plugin perangkat

Untuk menggunakan perangkat tertentu, seperti perangkat asli di TensorFlow, pengguna hanya perlu menginstal paket plug-in perangkat untuk perangkat tersebut. Cuplikan kode berikut menunjukkan bagaimana plugin untuk perangkat demonstrasi baru, Awesome Processing Unit (APU) , diinstal dan digunakan. Untuk mempermudah, contoh plugin APU ini hanya memiliki satu kernel khusus untuk 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

Dengan plug-in terpasang, uji apakah perangkat terlihat dan jalankan operasi pada perangkat APU baru:

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.

Perangkat yang tersedia

Perangkat Metal PluggableDevice untuk GPU macOS:

DirectML PluggableDevice untuk Windows dan WSL (pratinjau):

Ekstensi Intel® untuk TensorFlow PluggableDevice untuk Linux dan WSL: