Looping over tf.data.Dataset very slow

user_6920649 0
Viewed 1249

I'm wondering why a for-loop over samples of a tf.data.Dataset is so much slower than looping over the corresponding numpy array. 

import numpy as np
import tensorflow as tf
import time

a = np.ones(100000, dtype=np.float32)

start_time = time.time()
for x in a:
    pass
print(time.time() - start_time)

start_time = time.time()
for x in tf.data.Dataset.from_tensor_slices(a):
    pass
print(time.time() - start_time)

0.05548405647277832
5.67711615562439

My TensorFlow version is 2.0.0.

tensorflow-datasets tensorflow2.0 tensorflow python
1 Answers

Yes, even i have observed same behavior. To improve speed/performance try wrapping tf.data.dataset in a @tf.function and it will take almost the same time.

AutoGraph is on default in tf.function and transforms your Python eager code into graph-compatible TensorFlow ops. This includes control flow like if, for, while.

tf.function works best with TensorFlow ops, NumPy and Python calls are converted to constants.

Please refer code shown below to wrap within @tf.function

@tf.function

def oper(a):
    start_time = time.time()
    for x in tf.data.Dataset.from_tensor_slices(a):
        pass
    print(time.time() - start_time)

Complete working code shown below between numpy and tf.data.dataset performance

import numpy as np
import tensorflow as tf
import time

a = np.ones(100000, dtype=np.float32)

start_time = time.time()
for x in a:
    pass
print(time.time() - start_time)


@tf.function

def oper(a):
    start_time = time.time()

    for x in tf.data.Dataset.from_tensor_slices(a):
        pass
    print(time.time() - start_time)

oper(a)

output:

0.012496232986450195
0.017792224884033203

To know more about tf.function please refer this.

user_142906810
Related