为什么 tf.matmul 在 cpu 上运行

Question

I encountered some problems while training my model, specifically, the training speed is sometimes fast and sometimes slow(fluctuating between 2-10 iters per second). After I used tensorboard to check the.network operation, I found that the tf.matmul() node was running on CPU device. As follow:

This makes me quite confused, I try to use tf.device('/gpu:0') to enforce that node running on GPU, I also tried to lower the tensorflow version, but none of these work. below is my code:

def _repeat(x, n_repeats):
        with tf.variable_scope('_repeat'):
            rep = tf.transpose(
                tf.expand_dims(tf.ones(shape=tf_v.stack([n_repeats, ])), 1), [1, 0])
            rep = tf.cast(rep, tf.int32)
#             with tf.device('/gpu:0'):
            x = tf.matmul(tf.reshape(x, (-1, 1)), rep)
            return tf.reshape(x, [-1])

by the way, the tensorflow version is 2.2.5, GPU is rtx3090.

Answer 1

The possible reason could be that the GPU has not been configured in your system properly. This is why it's detecting the CPU only. Make sure you have Tensorflow GPU installed in your system along with other GPU support requirements.

If a TensorFlow operation has both CPU and GPU implementations, by default, the GPU device is prioritized when the operation is assigned.

I tried replicating tf.matmul() in Google colab by selecting "CPU mode" and it could not detect the GPU:

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
print(tf.config.list_physical_devices())

# Place tensors on the GPU
with tf.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]])
  c = tf.matmul(a, b)
  print(c)

Output:

Num GPUs Available:  0
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

But if we try the same by selecting "GPU mode", the tensor operation detects the GPU.

Output:

Num GPUs Available:  1
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Please check this link to get more details in this.