Plugin per dispositivi GPU

L'architettura dei dispositivi collegabili di TensorFlow aggiunge il supporto di nuovi dispositivi come pacchetti plug-in separati installati insieme al pacchetto TensorFlow ufficiale.

Il meccanismo non richiede modifiche specifiche del dispositivo nel codice TensorFlow. Si basa sulle API C per comunicare con il binario TensorFlow in modo stabile. Gli sviluppatori di plug-in mantengono repository di codice e pacchetti di distribuzione separati per i propri plug-in e sono responsabili del test dei propri dispositivi.

Utilizza i plugin del dispositivo

Per utilizzare un particolare dispositivo, come si farebbe con un dispositivo nativo in TensorFlow, gli utenti devono solo installare il pacchetto plug-in del dispositivo per quel dispositivo. Il seguente frammento di codice mostra come viene installato e utilizzato il plug-in per un nuovo dispositivo dimostrativo, Awesome Processing Unit (APU) . Per semplicità, questo plug-in APU di esempio ha solo un kernel personalizzato per 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

Con il plug-in installato, verifica che il dispositivo sia visibile ed esegui un'operazione sul nuovo dispositivo 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.

Dispositivi disponibili

Metal PluggableDevice per GPU macOS:

DirectML PluggableDevice per Windows e WSL (anteprima):

Estensione Intel® per TensorFlow PluggableDevice per Linux e WSL: