Tensorflow 2 eager vs graph mode

Question

I've been working through the tensorflow-2.0.0 beta tutorials. In the advanced example a tensorflow.keras subclass is used. The presence of the @tf.function decorator on train_step and test_step means the model executes in graph mode (not sure if that's the correct terminology, I mean oposite to eager mode). If I remove these decorators I can single step right into the model call function and see the input/output tensor for each layer which is neat.

My question is, is there a programatic way to enable/disable the @tf.function decorators. Commenting them out to switch between eager and graph mode doesn't seem particularly scaleable but it's certainly useful for debugging/learning)

Answer 1

You could always write two functions (one with the decorator and one without) and call whichever suits you...

For example

@tf.function()
def graph_function()
    # This function will operate in graph mode
    ...

def eager_function()
    # This function will operate in eager mode
   ...

if tf.executing_eagerly()
    my_function = eager_function
else:
    my_function = graph_function

# You proceed to my_function from now on

I don't know if there is a better way but I've seen this a lot being used in the tensorflow official repository on github.

Answer 2

I do not know in which version of Tensorflow it was introduced, but at least in TF 2.1, there is

tf.config.experimental_run_functions_eagerly(True)

available. It makes all @tf.function-decorated functions run in eager mode anyway, until this is reset by calling with argument False again.

For details see Tensorflow experimental_run_functions_eagerly documentation.