How can my Python Tensorflow program print what device { CPU, GPU, TPU } it is using?

user_2417922 0
Viewed 336

I'm trying to run a Deep Neural Net segmentation program, written in Python/Tensorflow running under Kaggle. I'm trying to understand how to control the device used in background runs [SAVE VERSION]. The first step is being able to see what device(s) I am running on -- CPU, GPU or TPU. How can my Python program determine which device(s) are currently in use?

kaggle tensorflow2.0 python-3.x python
2 Answers

There are many ways to answer your question. The simplest is to check what devices are available to you:

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

If you want to know some more details about each of the devices you can run tf.test.gpu_device_name to get the name of GPU device (or any other allocated).

user_65739020

Found a solution:

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 )
user_24179220
Related