Plug-ins de périphérique GPU

L'architecture des appareils enfichables de TensorFlow ajoute une nouvelle prise en charge des appareils sous forme de packages de plug-ins distincts installés avec le package TensorFlow officiel.

Le mécanisme ne nécessite aucune modification spécifique à l'appareil dans le code TensorFlow. Il s'appuie sur les API C pour communiquer avec le binaire TensorFlow de manière stable. Les développeurs de plug-ins gèrent des référentiels de code et des packages de distribution distincts pour leurs plug-ins et sont responsables du test de leurs appareils.

Utiliser des plugins d'appareil

Pour utiliser un appareil particulier, comme on le ferait avec un appareil natif dans TensorFlow, les utilisateurs n'ont qu'à installer le package de plug-in de périphérique pour cet appareil. L'extrait de code suivant montre comment le plugin pour un nouveau périphérique de démonstration, Awesome Processing Unit (APU) , est installé et utilisé. Par souci de simplicité, cet exemple de plug-in APU n'a qu'un seul noyau personnalisé pour 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

Une fois le plug-in installé, testez que le périphérique est visible et exécutez une opération sur le nouveau périphérique 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.

Appareils disponibles

Metal PluggableDevice pour les GPU macOS :

DirectML PluggableDevice pour Windows et WSL (aperçu) :

Extension Intel® pour TensorFlow PluggableDevice pour Linux et WSL :