CODEWITHSUNDEEP
tensorflow
JuhoKupiainen
JuhoKupiainen

Reputation: 109

How can I find out what device an op has been assigned to in Tensorflow?

I'm writing distributed Tensorflow code using Jupyter Notebook. I would like to make sure that my operations are placed on the correct devices. I would somehow like to inspect what operations have been assigned to what devices. I know this is possible using Tensorboard, but I was hoping there was a way to do this using Python.

Upvotes: 3

Views: 2172

Answers (2)

Mark Peschel
Mark Peschel

Reputation: 334

For Tensorflow 2, to log what devices are being used, use tf.debugging.set_log_device_placement(True).

For example, with the following snippet,

import tensorflow as tf
tf.debugging.set_log_device_placement(True)

print(tf.constant([[1, 2], [3, 4]]))

you should get output like

2023-06-02 21:03:30.359187: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.

...

input: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-02 21:03:33.562777: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
_EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-02 21:03:33.562811: I tensorflow/core/common_runtime/placer.cc:114] _EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:CPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-02 21:03:33.562872: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-02 21:03:33.575140: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:CPU:0
tf.Tensor(
[[1 2]
 [3 4]], shape=(2, 2), dtype=int32)

where the device in use (in this case, CPU:0) is printed at the end of each line.

Link to documentation

Upvotes: 1

JuhoKupiainen
JuhoKupiainen

Reputation: 109

Citing straight from Tensorflow documentation:

To find out which devices your operations and tensors are assigned to, create the session with log_device_placement configuration option set to True.

# Creates a graph.
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print(sess.run(c))

You should see the following output:

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K40c, pci bus
id: 0000:05:00.0
b: /job:localhost/replica:0/task:0/gpu:0
a: /job:localhost/replica:0/task:0/gpu:0
MatMul: /job:localhost/replica:0/task:0/gpu:0
[[ 22.  28.]
 [ 49.  64.]]

Upvotes: 3

Related Questions