如何检查我在 TensorFlow 中使用的是 CPU 还是 GPU？

Question

我读过

os.environ["CUDA_VISIBLE_DEVICES"] = ''

注意 tensorflow 将在 CPU 上运行，并且

os.environ["CUDA_VISIBLE_DEVICES"] = '0'

注意 tensorflow 将在 GPU 0 上运行。

我如何检查，使用的是哪个设备？

代码

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

只显示结果

[[ 22.  28.]
 [ 49.  64.]]

并且没有使用过的设备等。

Answer 1

您应该能够通过打开 tensorflow 日志记录语句来做到这一点。有几种方法可以做到这一点。您可以使用 bash 环境变量来做到这一点。

export TF_CPP_MIN_LOG_LEVEL=1

或在您的代码中使用..

tf.logging.set_verbosity(tf.logging.INFO)

在我的系统上，我得到类似...

设备映射：/job:localhost/replica:0/task:0/device:XLA_CPU:0 -> 设备：XLA_CPU 设备 /job:localhost/replica:0/task:0/device:XLA_GPU:0 -> 设备：XLA_GPU 设备 /job:localhost/replica:0/task:0/device:GPU:0 -> 设备: 0, 名称：GeForce GTX TITAN X，pci 总线 ID：0000:65:00.0，计算 能力：5.2 MatMul：（MatMul）： /job:localhost/replica:0/task:0/device:GPU:0 a: (Const): /job:localhost/replica:0/task:0/device:GPU:0 b:（常量）： /job:localhost/replica:0/task:0/device:GPU:0 [[22. 28.] [49. 64.]]

Answer 2

我发现 Python 和 IPython 之间存在差异。 IPython 是 Spyder 中使用的内核。所以我猜这就是导致这些不同输出的原因。