3
$\begingroup$

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)

$\endgroup$
4
$\begingroup$

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.

$\endgroup$
3
  • 1
    $\begingroup$ Thanks, I guess you could make it a bit cleaner by creating a factory method which returns either function. $\endgroup$ Aug 2 '19 at 6:33
  • $\begingroup$ Yeah, true. Tensorflow typically has different modules altogether for eager and graph modes, but I think that's mostly because different people work on each. $\endgroup$
    – Djib2011
    Aug 2 '19 at 6:52
  • $\begingroup$ I just found a recomendation to not decorate every function with @tf.function, rather just for example decorate a train_one_step and it'll inherit. So I guess that's not so bad, can create a debug and normal version $\endgroup$ Aug 2 '19 at 7:18
2
$\begingroup$

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.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.