python - 我怎么知道 tensorflow 张量是在 cuda 还是 cpu 中？

Question

我怎么知道tensorflow张量是在cuda还是cpu中？以这个非常简单的例子为例:

import tensorflow as tf
tf.debugging.set_log_device_placement(True)

# Place tensors on the CPU
with tf.device('/device:GPU:0'):
   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]])

# print tensor a
print(a)

# Run on the GPU
c = tf.matmul(a, b)
print(c)

代码运行良好。在这里，我将张量 'a' 和 'b' 物理放置在 GPU 上。在打印“a”时，我得到:

tf.Tensor(
  [[1. 2. 3.]
  [4. 5. 6.]], shape=(2, 3), dtype=float32)

它不提供任何信息，无论是 CPU 还是 GPU 中的“a”。现在，假设有一个中间张量，如在某些操作期间创建的张量 'c'。我怎么知道张量“c”是 CPU 还是 GPU 张量？
另外，假设张量放置在 GPU 上。我怎样才能把它移到 CPU 上？

Answer 1

从 Tensorflow 2.3 开始，您可以使用 .device张量的性质:

import tensorflow as tf

a = tf.constant([1, 2, 3])
print(a.device)  # /job:localhost/replica:0/task:0/device:CPU:0

更详细的解释可以找到 here