为什么Tensoflow 2.0和tf.function默认在CPU上处理某些功能

Question

我将tf.function与GPU版本的tensorflow一起使用。但是，对于某些功能，我必须明确地将GPU声明为用于张量流的设备（with tf.device('/device:GPU:0'):），否则这些功能将在CPU上执行。

例如，默认情况下，仅在一个CPU内核上执行将k稀疏性引入给定张量的功能：

@tf.function
def introdcue_k_sparsity(self, inputs; k):
    k_real = round(k * inputs.shape[-1])
    if k_real == 0:
        k_real = 1

    with tf.name_scope("sparsity_within_filter"):
        # Reorder pixel
        winners, _ = tf.nn.top_k(inputs, k_real)

        # Get all pixels that were sorted out and multiply with 0
        winners_inverse_shape, _ = tf.nn.top_k(inputs, inputs.shape[-1]-k_real)
        zeros = tf.scalar_mul(0, winners_inverse_shape)

        k_sparse_outputs = tf.concat([winners, zeros], axis=-1)

    return k_sparse_outputs