Complementos de dispositivo GPU

La arquitectura de dispositivos conectables de TensorFlow agrega compatibilidad con nuevos dispositivos como paquetes de complementos independientes que se instalan junto con el paquete oficial de TensorFlow.

El mecanismo no requiere cambios específicos del dispositivo en el código de TensorFlow. Se basa en las API de C para comunicarse con el binario de TensorFlow de manera estable. Los desarrolladores de complementos mantienen repositorios de código y paquetes de distribución separados para sus complementos y son responsables de probar sus dispositivos.

Usar complementos de dispositivo

Para usar un dispositivo en particular, como lo haría con un dispositivo nativo en TensorFlow, los usuarios solo tienen que instalar el paquete de complemento del dispositivo para ese dispositivo. El siguiente fragmento de código muestra cómo se instala y utiliza el complemento para un nuevo dispositivo de demostración, Awesome Processing Unit (APU) . Para simplificar, este complemento APU de muestra solo tiene un kernel personalizado para 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 el complemento instalado, pruebe que el dispositivo esté visible y ejecute una operación en el nuevo 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.

Dispositivos disponibles

PluggableDevice de metal para GPU macOS:

DirectML PluggableDevice para Windows y WSL (versión preliminar):

Extensión Intel® para TensorFlow PluggableDevice para Linux y WSL: