Wtyczki do urządzeń GPU

Architektura urządzeń wtykowych TensorFlow dodaje obsługę nowych urządzeń w postaci oddzielnych pakietów wtyczek, które są instalowane wraz z oficjalnym pakietem TensorFlow.

Mechanizm nie wymaga żadnych zmian specyficznych dla urządzenia w kodzie TensorFlow. Opiera się na interfejsach API języka C do stabilnej komunikacji z plikiem binarnym TensorFlow. Twórcy wtyczek utrzymują oddzielne repozytoria kodu i pakiety dystrybucyjne dla swoich wtyczek oraz są odpowiedzialni za testowanie swoich urządzeń.

Użyj wtyczek do urządzeń

Aby korzystać z określonego urządzenia, tak jak z urządzenia natywnego w TensorFlow, użytkownicy muszą jedynie zainstalować pakiet wtyczek urządzenia dla tego urządzenia. Poniższy fragment kodu pokazuje, jak instaluje się i używa wtyczki nowego urządzenia demonstracyjnego Awesome Processing Unit (APU) . Dla uproszczenia ta przykładowa wtyczka APU ma tylko jedno niestandardowe jądro dla 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

Po zainstalowaniu wtyczki sprawdź, czy urządzenie jest widoczne i wykonaj operację na nowym urządzeniu 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.

Dostępne urządzenia

Metalowe PluggableDevice dla procesorów graficznych macOS:

PluggableDevice DirectML dla systemu Windows i WSL (wersja zapoznawcza):

Rozszerzenie Intel® dla PluggableDevice TensorFlow dla systemu Linux i WSL: