我的 Python Tensorflow 程序如何打印它正在使用什么設備 { CPU, GPU, TPU }？

Question

我正在嘗試運行一個用 Python/Tensorflow 在 Kaggle 下運行的深度神經網絡分割程序。 我正在嘗試了解如何控制后台運行中使用的設備 [SAVE VERSION]。 第一步是能夠查看我正在運行的設備——CPU、GPU 或 TPU。 我的 Python 程序如何確定當前正在使用哪些設備？

Answer 1

有很多方法可以回答您的問題。 最簡單的方法是檢查您可以使用哪些設備：

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

如果您想了解有關每個設備的更多詳細信息，可以運行tf.test.gpu_device_name以獲取 GPU 設備的名稱（或任何其他分配的）。

Answer 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 )