py_function is slower in TF2.7 (and 2.6) compared to TF2.5

[anth2o]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• TensorFlow installed from (source or binary): docker
• TensorFlow version (use command below): 2.7
• Python version: 3.8.10

Describe the current behavior

In my training code I'm using tf.data.Dataset.map with a tf.py_function within it. I know that using tf.py_function is slower than regular tensorflow, but there's a part of my dataset processing that requires the usage of plain Python code.

However, when I tried to upgrade TF from 2.5 to 2.6, I noticed that my trainings were slower in 2.6 compared to 2.5. I postponed the upgrade at the moment. When 2.7 came out I tried again but the slower data processing was still present.

Describe the expected behavior

I would have expected this kind of code to be as efficient as it was in 2.5

Standalone code to reproduce the issue

You can find a reproducible example in this colab.

To reproduce my issue, I wrapped the opening of the image on disk within a tf.py_function and opened it using Pillow and numpy. I know that this can be done in a more efficient way as I explained directly in the colab.

Here are the results I got when executing the %%timeit cell with 2 different versions of TF:

• tf.version = "2.7.0" -> 3 loops, best of 7: 1.81 s per loop (would have been the same in TF2.6)
• tf.version = "2.5.2" -> 3 loops, best of 7: 1.44 s per loop

[mohantym]

Hi @chunduriv ! Could you please look at this issue? Attaching Gist in 2.7 and 2.6 for reference .Thank you!

Updated
Was able to replicate the issue with Tf-nightly-2.11.0.dev20220822 .Please find the gist here. Thank you!

[bgdax]

Thanks for raising this issue. We are looking into it.