Tensorflow中的可插拔处理单元

t98cgbkg  于 2023-06-24  发布在  其他
关注(0)|答案(1)|浏览(99)

他们有这样的例子:

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.

我的问题很简单:如何区分哪些操作在哪些设备上运行?在本例中,a在CPU中运行,b在可插拔设备上运行。
Source

brvekthn

brvekthn1#

在您引用的示例中,一个虚构的“Awesome Processing Unit”(APU)提供了ReLU操作(“* 为简单起见,此示例APU插件只有一个ReLU的自定义内核 *”)安装:

# Install the APU example plug-in package
$ pip install tensorflow-apu-0.0.1-cp36-cp36m-linux_x86_64.whl

报价:

  • 要使用特定设备,就像TensorFlow中的原生设备一样,用户只需安装该设备的设备插件包。

安装APU插件包后,APU设备可以在Tensorflow中作为本机设备使用。
可插拔设备具有自己的自定义操作和内核(参见https://blog.tensorflow.org/2021/06/pluggabledevice-device-plugins-for-TensorFlow.html)。来自同一个博客:

  • PluggableDevice type:TensorFlow中的新设备类型,允许从插件包注册设备。在器件放置阶段,它优先于本机器件。*

在你引用的例子中,引用:“*APU插件只有一个用于ReLU的自定义内核 *”,因此对于ReLU,默认情况下将选择APU,并且

c = tf.nn.relu(a)  # Runs on APU

相当于

with tf.device("/APU:0"):  
    c = tf.nn.relu(a)  # Runs on APU

我不知道如何列出插件设备注册的操作,也不知道在多个插件设备的情况下设备的优先级如何。

相关问题