TensorFlow中如何判断操作运行在哪个Pluggable Processing Unit设备？

下面是几种实用的方式，可直接确定TensorFlow中每个操作的运行设备：

1. 查看输出张量的.device属性 #

TensorFlow中，操作的运行设备通常与它输出的张量所在设备一致。你可以直接访问张量的device属性获取设备信息：

import tensorflow as tf

# 执行示例操作
a = tf.random.normal(shape=[5], dtype=tf.float32)
b = tf.nn.relu(a)

# 打印设备信息
print(f"生成张量a的操作运行设备：{a.device}")
print(f"计算张量b的relu操作运行设备：{b.device}")

输出示例：

生成张量a的操作运行设备：/job:localhost/replica:0/task:0/device:CPU:0
计算张量b的relu操作运行设备：/job:localhost/replica:0/task:0/device:APU:0

2. 启用设备分配日志 #

通过tf.debugging.set_log_device_placement(True)开启日志后，TensorFlow会在控制台打印每个操作的设备分配情况：

import tensorflow as tf

# 开启设备分配日志
tf.debugging.set_log_device_placement(True)

# 执行操作
a = tf.random.normal(shape=[5], dtype=tf.float32)
b = tf.nn.relu(a)

控制台会输出类似内容：

random_normal: (RandomNormal): /job:localhost/replica:0/task:0/device:CPU:0
relu: (Relu): /job:localhost/replica:0/task:0/device:APU:0

3. 导出调试信息到文件 #

对于复杂模型，可使用tf.debugging.experimental.enable_dump_debug_info将操作的设备信息导出到文件，方便后续查看：

import tensorflow as tf

# 启用调试信息导出，指定输出目录
tf.debugging.experimental.enable_dump_debug_info(
    "./tf_device_dump",
    tensor_debug_mode="NO_TENSORS",  # 仅保存结构和设备信息，不存储张量数据
    circular_buffer_size=-1
)

# 执行你的模型操作
a = tf.random.normal(shape=[5], dtype=tf.float32)
b = tf.nn.relu(a)

导出的文件中会包含每个操作的device字段，可直接查看对应设备。

4. 手动指定设备上下文验证 #

使用tf.device()上下文管理器强制指定操作运行的设备，对比默认分配结果，也能验证操作的实际运行设备：

import tensorflow as tf

# 强制在CPU上执行操作
with tf.device("/CPU:0"):
    c = tf.random.normal(shape=[5])
print(f"强制CPU运行的操作输出张量设备：{c.device}")

# 强制在APU上执行操作
with tf.device("/APU:0"):
    d = tf.nn.relu(c)
print(f"强制APU运行的relu操作输出张量设备：{d.device}")

内容的提问来源于stack exchange，提问作者user12582392