Python – tensorflow.device()

TensorFlow是谷歌设计的开源Python库，用于开发机器学习模型和深度学习神经网络。

device()是用来明确指定操作应该在哪个设备上进行的。

语法: tensorflow.device( device_name )

参数:

• device_name: 它指定了在此背景下使用的设备名称。

返回：它返回一个上下文管理器，指定用于新创建的操作的默认设备。

示例 1:

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

输出:

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

例子2：在这个例子中，设备规格指定使用GPU，但系统找不到GPU，所以它将在CPU上运行操作。

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

输出:

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