如何知道 TensorFlow Tensor 是在 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)

它没有提供任何’a‘在CPU或GPU中的信息。现在，假设有一个类似于张量’c‘的中间张量，在某些操作期间创建。我如何知道张量’c‘是CPU还是GPU张量？ 另外，假设张量被放置在GPU上。我该如何将它移动到CPU？

- Kumar Govindam

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

更详细的说明可以在此处找到。

- sebastian-sz

Answer 2

你可以参考PyTorch中的内存管理，其中你可以明确地定义张量保存在哪个内存中。据我所知，这在Tensorflow（2.X版本）中并不支持，你只能在CPU或GPU上工作。这是在声明张量时根据你的TF版本决定的。据我所知，默认情况下会使用GPU，否则必须在开始任何图形操作之前明确指定。

经验法则：如果你有一个可用的cuda环境和一个默认支持GPU的TF版本，它将始终运行在GPU上，否则在CPU上，除非你手动定义。

参考Patwie在SO上的answer。

- MichaelJanz

---

正如你提到的，在pytorch中的内存管理中，可以选择在CPU/GPU中定义张量。当我在研究一些pytorch模型库时，代码明确将一些数据从GPU移动到CPU，因为特定操作在CPU上更快。我正在寻找这种功能，但我想这并不存在。 - Kumar Govindam

抱歉，不能。不过，可以使用 eager tensor 上的 .numpy() 方法来帮助您。它会返回一个 numpy 数组，然后存储在 CPU 内存中。 - MichaelJanz