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

Question

我在python 3.8中使用Tensorflow 2.2。我有一个张量切片的数据集对象构建，需要在数据集的每个张量上应用一些计算，称之为 。为此，我使用了（请参阅下面的代码）的功能。然而，与在每个张量上直接应用给定方法相比，该映射的执行速度相当慢。这是模型案例（下面的代码保存在名为 的文件中）。computemaptf.data.Datasettest.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和直接应用用于地图的相同函数在性能上存在如此巨大的差距？

当我相应地添加设置num_parallel_calls的deterministic参数（我的机器上的核心数）并且进程运行时（与没有并行化相比）。尽管如此，这仍然比直接应用 中使用的方法要差得多。map8False~0.16 sec~1 secmap

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

Answer 1

我正在回答我的问题，以防万一有人遇到问题中描述的相同问题。以下内容基于此Github 问题的评论（其中提供了更多信息）。

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

tf.debugging.set_log_device_placement(True) 

到代码以访问有关 Tensorflow 操作位置的信息。要强制map在 GPU 上执行，可以compute在

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

块（参见上面的链接）。