Search code examples
pythontensorflowtensorflow-estimator

Eager execution in tensorflow custom estimators

I'm rewriting my code using custom Estimator logic, and I need to enable eager execution to get the metrics/predictions I need. However, it seems that enabling eager execution does not go through for some reason. To reproduce, I can use the example located at https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/learn/iris.py with some prints:

...
import tensorflow as tf
tf.enable_eager_execution()
...
def my_model(features, labels, mode):
  print("IS EAGER? (my_model) - {}".format(tf.executing_eagerly()))
...
print("IS EAGER? - {}".format(tf.executing_eagerly()))
classifier = tf.estimator.Estimator(model_fn=my_model)

Which results in the following when I run the script:

IS EAGER? - True
INFO:tensorflow:Using default config.
...
INFO:tensorflow:Calling model_fn.
IS EAGER? (my_model) - False
INFO:tensorflow:Done calling model_fn.

How do I get my model to execute eagerly? I'm using tensorflow 1.9.0

Solution

  • The Estimator API is pretty tightly tied to graph construction (each call to train(), evaluate() etc. reconstructs the graph). So it explicitly disables eager execution when invoked.

    Is it possible for you to use your own training loop or use tf.keras.Model.fit() instead of using the Estimator API?