ipython.out

In [7]: import numpy as np
   ...: import tensorflow as tf
   ...: print(tf.version.GIT_VERSION, tf.version.VERSION)
   ...: 
   ...: dim = 300
   ...: 
   ...: def return_single(*args):
   ...:     x = np.random.rand(dim,dim,dim)
   ...:     return x,1
   ...: 
   ...: def return_list(*args):
   ...:     x = np.random.rand(dim,dim,dim)
   ...:     return [x], [1]
   ...: 
   ...: def return_nd(*args):
   ...:     x = np.random.rand(dim,dim,dim)
   ...:     return np.array([x]), np.array([1])
   ...: 
   ...: def make_ds(fun):
   ...:     the_ds = tf.data.Dataset.range(5)
   ...:     the_ds = the_ds.map(
   ...:                 lambda fn: tf.py_function(
   ...:                     fun, [fn],
   ...:                     [tf.float32,tf.float32]
   ...:                 )
   ...:             )
   ...:     return the_ds
   ...: 
v2.4.0-49-g85c8b2a817f 2.4.1

In [8]: %%time
   ...: for _ in make_ds(return_single):
   ...:     pass
   ...: 
CPU times: user 1.02 s, sys: 435 ms, total: 1.46 s
Wall time: 1.46 s

In [9]: %%time
   ...: for _ in make_ds(return_nd):
   ...:     pass
   ...: 
CPU times: user 1.2 s, sys: 850 ms, total: 2.05 s
Wall time: 2.06 s

In [10]: %%time
    ...: for _ in make_ds(return_list):
    ...:     pass
    ...: 
CPU times: user 11.4 s, sys: 345 ms, total: 11.7 s
Wall time: 11.7 s

tf_ds_py_function_performance_glitch.py

import numpy as np
import tensorflow as tf
print(tf.version.GIT_VERSION, tf.version.VERSION)

dim = 300

def return_single(*args):
    x = np.random.rand(dim,dim,dim)
    return x,1

def return_list(*args):
    x = np.random.rand(dim,dim,dim)
    return [x], [1]

def return_nd(*args):
    x = np.random.rand(dim,dim,dim)
    return np.array([x]), np.array([1])

def make_ds(fun):
    the_ds = tf.data.Dataset.range(5)    
    the_ds = the_ds.map(
                lambda fn: tf.py_function(
                    fun, [fn], 
                    [tf.float32,tf.float32]
                )
            )
    return the_ds
    
    
    
%%time
for _ in make_ds(return_single):
    pass
  
%%time
for _ in make_ds(return_nd):
    pass
  
%%time
for _ in make_ds(return_list):
    pass