Python – tensorflow.device()

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

device()用于明确指定应执行操作的设备。

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. 

示例 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)

输出：

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 上运行操作。

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)

输出：

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