Plugin thiết bị GPU

Kiến trúc thiết bị có thể cắm của TensorFlow bổ sung hỗ trợ thiết bị mới dưới dạng các gói plug-in riêng biệt được cài đặt cùng với gói TensorFlow chính thức.

Cơ chế này không yêu cầu thay đổi dành riêng cho thiết bị nào trong mã TensorFlow. Nó dựa vào API C để giao tiếp với nhị phân TensorFlow một cách ổn định. Các nhà phát triển plugin duy trì các kho mã và gói phân phối riêng biệt cho plugin của họ và chịu trách nhiệm kiểm tra thiết bị của họ.

Sử dụng plugin thiết bị

Để sử dụng một thiết bị cụ thể, giống như một thiết bị gốc trong TensorFlow, người dùng chỉ phải cài đặt gói plug-in thiết bị cho thiết bị đó. Đoạn mã sau đây cho biết cách cài đặt và sử dụng plugin cho thiết bị trình diễn mới, Bộ xử lý tuyệt vời (APU) . Để đơn giản, plug-in APU mẫu này chỉ có một kernel tùy chỉnh cho 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

Sau khi cài đặt plug-in, hãy kiểm tra xem thiết bị có hiển thị không và chạy thao tác trên thiết bị APU mới:

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.

Thiết bị có sẵn

PluggableDevice kim loại dành cho GPU macOS:

DirectML PluggableDevice dành cho Windows và WSL (xem trước):

Tiện ích mở rộng Intel® dành cho TensorFlow PluggableDevice dành cho Linux và WSL: