在张量流数据集上应用地图执行非常缓慢

Question

我在 python 3.8 中使用 Tensorflow 2.2。我有一个张量切片的数据集对象构建，需要在数据集的每个张量上应用一些计算，称之为compute。为此，我使用了tf.data.Dataset 的map 功能（代码见下文）。然而，与在每个张量上直接应用给定方法相比，该映射的执行速度相当慢。这是模型案例（下面的代码保存在一个名为test.py的文件中）。

import tensorflow as tf

class Test:
    def __init__(self):
        pass

    @tf.function
    def compute(self, tensor):
        # the main function that performs some computation with a tensor
        print('python.print ===> tracing compute ... ')

        res = tensor*tensor
        res = tf.signal.rfft(res)  # perform some computationally heavy task

        return res

    def apply_on_ds(self, ds):
        # mapping the compute method on a dataset
        return ds.map(lambda x: self.compute( x ) )

    @tf.function
    def apply_on_tensors(self, tensors):
        # a direct application on tensors of the compute method
        for i in tf.range(tensors.shape[0]):
            res = self.compute(tensors[i] )

要运行上面存储在test.py 中的代码，我执行以下操作

import tensorflow as tf
import time

import test

T = test.Test()
tensors = tf.random.uniform(shape=[100, 10000], dtype=tf.float32)
ds      = tf.data.Dataset.from_tensor_slices(tensors)

t1 = time.time(); x = list( T.apply_on_ds(ds) );  t2 = time.time();
# t2 - t1 equals ~1.08 sec on my computer

t1 = time.time() ;  x = T.apply_on_tensors(tensors);   t2 = time.time();
# t2 - t1 equals ~0.03 sec on my computer

为什么在应用 map 并直接应用与地图相同的功能？

当我将map 设置的num_parallel_calls 和deterministic 参数添加到相应的8（我机器上的内核数）和False 时，进程在~0.16 sec 中运行（与~1 sec 相比）没有并行化）。尽管如此，这仍然比直接应用map中使用的方法差很多。

我在这里有什么明显的错误吗？我怀疑在使用地图时对图表进行了一些回溯，但是，我找不到这方面的证据。对上述任何解释和改进建议将不胜感激。

Answer 1

我正在回答我的问题，以防万一有人遇到问题中描述的相同问题。以下内容基于 Github issue 上的 cmets（更多信息可在其中获得）。

代码本身没有问题。性能上的差距是因为map 操作（op）是在CPU 上放置和执行的，而map 中使用的函数的一对一应用是在GPU，因此性能差异。要看到这一点，可以添加

tf.debugging.set_log_device_placement(True) 

访问代码以访问有关 Tensorflow 将其操作放置在何处的信息。 要强制 map 在 GPU 上执行，可以在内部计算 compute 方法

with tf.device("/gpu:0"):    

阻止（参见上面的链接）。