Python – tensorflow.device()
Last Updated :
06 Aug, 2021
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
import tensorflow as tf
device_spec = tf.DeviceSpec(job = "localhost" , replica = 0 , device_type = "CPU" )
print ( 'Device Spec: ' , device_spec.to_string())
tf.debugging.set_log_device_placement( True )
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
import tensorflow as tf
device_spec = tf.DeviceSpec(job = "localhost" , replica = 0 , device_type = "GPU" )
print ( 'Device Spec: ' , device_spec.to_string())
tf.debugging.set_log_device_placement( True )
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
Like Article
Suggest improvement
Share your thoughts in the comments
Please Login to comment...