Applying map on tensorflow Dataset performs very slowly

Question

I work with Tensorflow 2.2 in python 3.8. I have a Dataset object build of tensor slices and need to apply some computation, call it compute, on each tensor of the dataset. To this end I use the map functionality of tf.data.Dataset (see below for the code). The map, however, performs rather slowly compared to direct application of the given method on each tensor. Here is the model case (the code below is saved in a file called 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] )

To run the code above, stored in test.py, I do the following

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

Why is there such a massive gap in performance between applying the map and applying the same function used for map directly ?

When I add num_parallel_calls and deterministic parameters of the map setting to correspondingly 8 (number of cores on my machine) and False the process runs in ~0.16 sec (vs. ~1 sec without parallelization). Nevertheless, this is still significantly worse than the direct application of the method used in map.

Is there any apparent mistake I am doing here? I suspect there is some retracing of the graph in place when using the map, however, I could not find an evidence for that. Any explanation of the above and suggestions for improvement would be much appreciated.

user_5747100 · Answer

I am answering my question, just in case someone might encounter the same problem described in the question. What follows is based on the comments on this Github issue (more information is available therein).

There is no problem with the code itself. The gap in performance is because the map operation (op) is placed and executed on CPU while the one-by-one application of the function used in the map is taking place on GPU, hence the difference in the performance. To see this, one can add

tf.debugging.set_log_device_placement(True)

to the code to access the information on where Tensorflow places its operations. To force the map to execute on GPU, one can take the computation of the compute method inside the

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

block (see the link above).

Applying map on tensorflow Dataset performs very slowly

1 Answers