ปลั๊กอินอุปกรณ์ GPU

สถาปัตยกรรม อุปกรณ์แบบเสียบได้ ของ TensorFlow เพิ่มการรองรับอุปกรณ์ใหม่เป็นแพ็คเกจปลั๊กอินแยกต่างหากที่ติดตั้งควบคู่ไปกับแพ็คเกจ TensorFlow อย่างเป็นทางการ

กลไกนี้ไม่จำเป็นต้องมีการเปลี่ยนแปลงเฉพาะอุปกรณ์ในโค้ด TensorFlow โดยอาศัย C API เพื่อสื่อสารกับไบนารี 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 GPU:

DirectML PluggableDevice สำหรับ Windows และ WSL (ตัวอย่าง):

ส่วนขยาย Intel® สำหรับ TensorFlow PluggableDevice สำหรับ Linux และ WSL: