Plug-ins de dispositivo GPU

A arquitetura de dispositivos conectáveis ​​do TensorFlow adiciona suporte a novos dispositivos como pacotes de plug-ins separados que são instalados junto com o pacote oficial do TensorFlow.

O mecanismo não requer alterações específicas do dispositivo no código do TensorFlow. Ele se baseia em APIs C para se comunicar com o binário TensorFlow de maneira estável. Os desenvolvedores de plug-in mantêm repositórios de código e pacotes de distribuição separados para seus plug-ins e são responsáveis ​​por testar seus dispositivos.

Usar plug-ins do dispositivo

Para usar um dispositivo específico, como um dispositivo nativo no TensorFlow, os usuários só precisam instalar o pacote de plug-in do dispositivo para esse dispositivo. O trecho de código a seguir mostra como o plug-in para um novo dispositivo de demonstração, Awesome Processing Unit (APU) , é instalado e usado. Para simplificar, este plug-in APU de amostra possui apenas um 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

Com o plug-in instalado, teste se o dispositivo está visível e execute uma operação no novo 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 disponíveis

Metal PluggableDevice para GPUs macOS:

DirectML PluggableDevice para Windows e WSL (visualização):

Extensão Intel® para TensorFlow PluggableDevice para Linux e WSL: