TensorFlow 2.0: tf.function and AutoGraph | by Michael Grogan | Towards Data Science

With the advent of TensorFlow (TF) 2.0, the introduction of tf.function has brought about some useful improvements to TF 1.0, most notably the introduction of AutoGraph.
Fundamentally, TensorFlow runs by means of computational graphs — i.e. a graph of nodes is used to represent a series of TensorFlow operations.
However, compiling TensorFlow graphs can be a cumbersome process, given that the syntax structure differs significantly from Python, and graph code uses more resources than a simple Python structure.
The purpose of AutoGraph is to graph code through the transformation of code written in Python’s classic syntax structure into TensorFlow graph-compatible code. This makes using Python and TensorFlow much more intuitive, as it eases the transition from Python code to graph code.
More fundamentally, tf.function allows for the intuitive use of both eager execution and AutoGraph, whereby a function can be run using Python syntax initially and then transferred into the equivalent graph code.
Let’s take an example.
Here is a Python function for calculating the area of a circle. Eager execution is turned on by default in TF 2.0, and the value r (the radius of a circle) is being defined as a tf.Variable.
>>> import tensorflow as tf>>> tf.executing_eagerly()>>> r = tf.Variable(10.0, name="r")>>> def area(r):>>>     circle=3.14*(r**2.00)>>>     return circle>>> area(10)314.0>>> print('Areas of respective circles: %2.2f, %2.2f, %2.2f' % (area(tf.constant(50.0)), area(tf.constant(100.0)), area(tf.constant(150.0))))Areas of respective circles: 7850.00, 31400.00, 70650.00
Now, suppose one wishes to implement this code using AutoGraph. Using TF 1.0 conventions, this would have been possible using tf.Session. Previously, a graph had to be explicitly launched in a Session in order to run.
>>> r = tf.Variable(10.0, name="r")>>> def area(r):>>>     circle=3.14*(r**2.00)>>>     return circle>>> print(tf.autograph.to_code(area))def tf__area(r):  do_return = False  retval_ = ag__.UndefinedReturnValue()  circle = 3.14 * r ** 2.0  do_return = True  retval_ = circle  cond = ag__.is_undefined_return(retval_)  def get_state():    return ()  def set_state(_):    pass  def if_true():    retval_ = None    return retval_  def if_false():    return retval_  retval_ = ag__.if_stmt(cond, if_true, if_false, get_state, set_state)  return retval_
This is a sample of the code in graph format.
>>> tf_area = tf.autograph.to_graph(area)>>> tf_area<function __main__.create_converted_entity_factory.<locals>.create_converted_entity.<locals>.tf__area(r)>>>> print(tf_area)<function create_converted_entity_factory.<locals>.create_converted_entity.<locals>.tf__area at 0x7f7b57f8a620>>>> with tf.Graph().as_default():    a1 = tf_area(tf.constant(50.0))    a2 = tf_area(tf.constant(100.0))    a3 = tf_area(tf.constant(150.0))    with tf.compat.v1.Session() as sess:        print('Areas of respective circles (graph results): %2.2f, %2.2f, %2.2f\n' % (sess.run(a1), sess.run(a2), sess.run(a3)))Areas of respective circles (graph results): 7850.00, 31400.00, 70650.00
The use of a Session means that the area values for the circles calculated above will only be computed inside the Session itself.
Let’s take a look at how this could be implemented using tf.function.
>>> r = tf.Variable(10.0, name="r")>>> @tf.function>>> def area(r):>>>     circle=3.14*(r**2.00)>>>     return circle>>> print(tf.autograph.to_code(area.python_function))def tf__area(r):  do_return = False  retval_ = ag__.UndefinedReturnValue()  circle = 3.14 * r ** 2.0  do_return = True  retval_ = circle  cond = ag__.is_undefined_return(retval_)  def get_state():    return ()  def set_state(_):    pass  def if_true():    retval_ = None    return retval_  def if_false():    return retval_  retval_ = ag__.if_stmt(cond, if_true, if_false, get_state, set_state)  return retval_>>> area(60)<tf.Tensor: id=13, shape=(), dtype=float32, numpy=11304.0>>>> print(area(tf.constant(50.00)), area(tf.constant(100.0)), area(tf.constant(150.0)))tf.Tensor(7850.0005, shape=(), dtype=float32) tf.Tensor(31400.002, shape=(), dtype=float32) tf.Tensor(70650.0, shape=(), dtype=float32)
As can be seen from the above, a Session is not needed to compute the area values. Instead, tf.function is used to specify the function which can be directly converted to AutoGraph, and then the area values are calculated as Tensors.
In this example, we have seen how TF 2.0 simplifies function implementation in TensorFlow through the use of tf.function and AutoGraph. Many thanks for your time, and you can also find more examples of machine learning using TensorFlow and Keras at michael-grogan.com.
Disclaimer: This article is written on an “as is” basis and without warranty. It was written with the intention of providing an overview of data science concepts, and should not be interpreted as professional advice in any way.