Плагины устройств GPU

Архитектура подключаемых устройств TensorFlow добавляет поддержку новых устройств в виде отдельных пакетов подключаемых модулей, которые устанавливаются вместе с официальным пакетом TensorFlow.

Этот механизм не требует никаких изменений кода TensorFlow для конкретного устройства. Он использует API C для стабильного взаимодействия с двоичным файлом TensorFlow. Разработчики плагинов поддерживают отдельные репозитории кода и пакеты распространения для своих плагинов и несут ответственность за тестирование своих устройств.

Используйте плагины устройств

Чтобы использовать конкретное устройство, например собственное устройство в TensorFlow, пользователям достаточно установить пакет подключаемого модуля устройства для этого устройства. В следующем фрагменте кода показано, как устанавливается и используется плагин для нового демонстрационного устройства Awesome Processing Unit (APU) . Для простоты этот пример плагина APU имеет только одно собственное ядро ​​для 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

Установив плагин, проверьте, отображается ли устройство, и запустите операцию на новом устройстве 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.

Доступные устройства

Металлическое PluggableDevice для графических процессоров macOS:

DirectML PluggableDevice для Windows и WSL (предварительная версия):

Расширение Intel® для TensorFlow PluggableDevice для Linux и WSL: