In tensorflow 2. what becomes a part of the graph and what does not?

user_8089695 0
Viewed 536

In the code:

@tf.function
def f(x):
    print("Python execution: ", x)
    out = x+5.
    tf.print("Graph execution: ", x)
    return out

python print executes only during the first invocation of the function i.e. during the graph creation. out=x+5. becomes a part of the graph and thus gets executed every time the graph is run. I wanted to know when i write a function in normal python and decorate it with tf.function to graph it up, what all gets converted to graph (like out = x+5) and what does not (like python print)

tensorflow2.x tensorflow2.0 tensorflow
2 Answers

Elaborating jdehesa's comment for the benefit of the community.

In the Official Documentation for tf.function, it is mentioned as:

Key Point: Any Python side-effects (appending to a list, printing with print, etc) will only happen once, when func is traced. To have side-effects executed into your tf.function they need to be written as TF ops:

Applying that explanation to your example, the Operations, out = x+5 and tf.print will be stored in the Graph and the operation print doesn't.

It is always recommended to pass Tensors as arguments to the functions instead of Python Scalars so that all the operations will be captured in a Single Graph.

tf.function only allows creating new tf.Variable objects when it is called for the first time.

In general, it is recommended to create stateful objects like tf.Variable outside of tf.function and passing them as arguments.

If we use Python Assignment Operator, a = b, will not be part of the Graph and tf.assign will be. Python Print will not be part of the Graph but tf.Print will be.

Also, if we want even Python Code to be part of Graph, we can use tf.py_function inside our function, but it is not recommended.

Hope this helps. Happy Learning!!

user_134652580

As the previous answer explains, only tf. calls (operations) are transformed by a tf.function.

To understand why, you should think what it means to write tf.print("msg") in a python program. It says:

  1. Define a Tensorflow operation (op for short) that does something. (It prints as side-effect, in this case, but the actual result is not important)
  2. Run. This happens because eager execution is enabled by default in TF2.

Essentially, when you decorate with tf.function, all the ops you define are collected, and reused later. So, every istruction that doesn't define a Tensorflow Op can't be recorded by Tensorflow. The + sign does define an op implicitly, in you example, because tf.function converts x to a Tensor, then pass it to f(x).

To answer to the most recent doubts:

  1. Yes, numpy arrays, and other native types, are treated as tensors, because TF ops and functions (a @tf.function decorated def) first call tf.convert_to_tensor on their inputs.
  2. No, a=b is not recorded because it simply copies the reference of a Tensor from one variable to the other. Assigments don't transform tensors.
user_67885710
Related