GPU 장치 플러그인

TensorFlow의 플러그형 장치 아키텍처는 공식 TensorFlow 패키지와 함께 설치되는 별도의 플러그인 패키지로 새로운 장치 지원을 추가합니다.

이 메커니즘은 TensorFlow 코드에서 장치별 변경이 필요하지 않습니다. C API를 사용하여 TensorFlow 바이너리와 안정적인 방식으로 통신합니다. 플러그인 개발자는 플러그인에 대한 별도의 코드 저장소와 배포 패키지를 유지 관리하고 장치 테스트를 담당합니다.

장치 플러그인 사용

TensorFlow의 기본 장치처럼 특정 장치를 사용하려면 사용자는 해당 장치에 대한 장치 플러그인 패키지만 설치하면 됩니다. 다음 코드 조각은 새로운 데모 장치인 APU(Awesome Process Unit) 용 플러그인이 설치되고 사용되는 방법을 보여줍니다. 단순화를 위해 이 샘플 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.

사용 가능한 장치

macOS GPU용 Metal PluggableDevice :

Windows 및 WSL용 DirectML PluggableDevice (미리 보기):

Linux 및 WSL용 TensorFlow PluggableDevice 용 Intel® 확장: