Can't enable eager execution mode in map function

[breadbread1984]

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

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): ubntu 18.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.0.0.dev20190412
• Python version: 3.6.7
• Bazel version (if compiling from source): N/A
• GCC/Compiler version (if compiling from source): N/A
• CUDA/cuDNN version: N/A
• GPU model and memory: N/A

You can collect some of this information using our environment capture script
You can also obtain the TensorFlow version with
python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"

Describe the current behavior

I write some preprocess code which can only be executed eagerly in map function for dataset formating, but I find that the code can't be executed properly. Even if the main function is executed in eager mode, the map function is still executed in graph mode.

Describe the expected behavior

I hope the map function can also work in eager mode to ease the coding difficulty.

Code to reproduce the issue
Provide a reproducible test case that is the bare minimum necessary to generate the problem.

assert tf.executing_eagerly can trigger an assertion failure.

Other info / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

[captain-pool]

@breadbread1984 Can you provide a small code snippet which can just be copy-pasted in the interpreter to reproduce the issue?

[breadbread1984]

@captain-pool this code can reproduce the problem.

#!/usr/bin/python3

import tensorflow as tf;
import tensorflow_datasets as tfds;

def map_func(feature):

    # assertion fails
    assert tf.executing_eagerly();
    return feature;

def main():

    trainset = tfds.load(name = "mnist", split = tfds.Split.TRAIN);
    trainset = trainset.map(map_func);
    for feature in trainset:
        pass;

if __name__ == "__main__":

    # assertion passes
    assert tf.executing_eagerly();
    main();

[wolffg]

In 2.0, code inside Datasets maps is turned into a subgraph for speed, just as it was in 1.x eager execution. You generally want to avoid Python inside your data pipeline. Otherwise, when you start using accelerators you could easily get I/O bound.

In the future for 2.0, tf.function (and thus Autograph) will automatically be applied to those functions, but it doesn't happen yet.

One easy way to run Python eagerly on your Dataset is to do it after it comes out of the Dataset within an eager loop (just call your transformation on each batch). If being inside the map function is key, you can use py_function to force Python execution. See:

https://www.tensorflow.org/alpha/tutorials/load_data/tf_records#writing_a_tfrecord_file

Various footnotes from 1.x Eager discussions:
#14732
https://stackoverflow.com/questions/50538038/tf-data-dataset-mapmap-func-with-eager-mode
https://stackoverflow.com/questions/49270477/tensorflow-dataset-map-api