获取设备名称 TensorFlow - Python

在 TensorFlow 中，我们可以使用 tf.device 函数来指定当前操作运行的设备，但是如何获取当前设备的名称呢？我们可以使用 tf.test.gpu_device_name() 来获取当前 GPU 设备名称，使用 tf.test.is_built_with_cuda() 来判断 TensorFlow 是否使用了 CUDA 加速。

import tensorflow as tf

if tf.test.is_built_with_cuda():
    print("TensorFlow 使用了 CUDA 加速")
    print("当前 GPU 设备名称：", tf.test.gpu_device_name())
else:
    print("TensorFlow 没有使用 CUDA 加速")

以上代码输出结果如下：

TensorFlow 使用了 CUDA 加速
当前 GPU 设备名称： /device:GPU:0

当然，如果我们使用的是 CPU 环境，则 tf.test.gpu_device_name() 会返回 None。

以上就是获取设备名称 TensorFlow - Python 的介绍。