如何确定 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?
.numpy()
方法来帮助您。它会返回一个 numpy 数组,然后存储在 CPU 内存中。 - MichaelJanz