我的 Python TensorFlow 程序如何打印它正在使用的设备 { CPU、GPU、TPU }？答案

【问题标题】：How can my Python Tensorflow program print what device { CPU, GPU, TPU } it is using?我的 Python TensorFlow 程序如何打印它正在使用的设备 { CPU、GPU、TPU }？
【发布时间】：2021-04-21 02:22:21
【问题描述】：

我正在尝试运行一个用 Python/Tensorflow 在 Kaggle 下运行的深度神经网络分割程序。我正在尝试了解如何控制后台运行中使用的设备 [SAVE VERSION]。第一步是能够查看我在哪些设备上运行——CPU、GPU 或 TPU。我的 Python 程序如何确定当前正在使用哪些设备？

【问题讨论】：

标签： python python-3.x tensorflow2.0 kaggle

【解决方案1】：

有很多方法可以回答您的问题。最简单的方法是检查您可以使用哪些设备：

with tf.Session() as sess:
  devices = sess.list_devices()

如果您想了解有关每个设备的更多详细信息，您可以运行 tf.test.gpu_device_name 以获取 GPU 设备的名称（或任何其他分配的）。

【讨论】：

谢谢。两个问题：首先，我不得不将 sess 的定义更改为tf.compat.v1.Session()。其次，虽然它在我启用 GPU 时正确列出了 GPU 设备，但在启用时不列出 TPU 设备。这是因为 TPU 是 Tensorflow 2.0 的功能吗？
另外，tf.config.list_physical_devices(device_type=None) 不显示 TPU。

【解决方案2】：

找到解决办法：

def running_on_TPU():
    try:
        tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection
        tf.config.experimental_connect_to_cluster(tpu)
        tf.tpu.experimental.initialize_tpu_system(tpu)
        strategy = tf.distribute.experimental.TPUStrategy(tpu)
        return True
    except:
        return False

print( "running_on_TPU", running_on_TPU(), file = sys.stderr )

【讨论】：