Python – tensorflow.device()

TensorFlow is open-source Python library designed by Google to develop Machine Learning models and deep learning  neural networks.

device() is used to explicitly specify the device in which operation should be performed.

Syntax: tensorflow.device( device_name )

Parameters:

• device_name: It specifies the device name to be used in this context.

Returns: It returns a context manager that specifies the default device to use for newly created ops. 

Example 1:

Python3

# Importing the library
import tensorflow as tf

# Initializing Device Specification
device_spec = tf.DeviceSpec(job ="localhost", replica = 0, device_type = "CPU")

# Printing the DeviceSpec
print('Device Spec: ', device_spec.to_string())

# Enabling device logging
tf.debugging.set_log_device_placement(True)

# Specifying the device
with tf.device(device_spec):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
  c = tf.matmul(a, b)

Output:

Device Spec:  /job:localhost/replica:0/device:CPU:*
Executing op MatMul in device /job:localhost/replica:0/task:0/device:CPU:0

Example 2: In this example device specification specifies GPU to be used but system couldn’t find GPU so it will run the operation on CPU. 

Python3

# Importing the library
import tensorflow as tf

# Initializing Device Specification
device_spec = tf.DeviceSpec(job ="localhost", replica = 0, device_type = "GPU")

# Printing the DeviceSpec
print('Device Spec: ', device_spec.to_string())

# Enabling device logging
tf.debugging.set_log_device_placement(True)

# Specifying the device
with tf.device(device_spec):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
  c = tf.matmul(a, b)

Output:

Device Spec:  /job:localhost/replica:0/device:GPU:*
Executing op MatMul in device /job:localhost/replica:0/task:0/device:CPU:0