The @tf.function decorator in TensorFlow is used to convert a
Python function into a TensorFlow graph function. This decorator
allows you to optimize and accelerate the execution of the
function by converting it into a TensorFlow graph, which can be
efficiently executed on GPUs or TPUs.
@tf.function
def tf_cube(x):
return x ** 3
result = tf_cube.python_function(2)
print(result)
The tf_cube function is decorated with @tf.function. This
indicates that the function should be compiled into a
TensorFlow graph. When the function is called, TensorFlow
traces the execution and constructs a graph representation of
the operations performed in the function.
The tf_cube.python_function(2) line invokes the TensorFlow
graph function for the tf_cube function with an input value of 2.
The graph function applies the cube operation (x ** 3) to the
input and returns the result.
By using the @tf.function decorator, you can take advantage of
the benefits of graph execution, such as automatic
differentiation, potential optimizations, and the ability to run on
accelerators. The TensorFlow graph function can be executed
efficiently multiple times with different inputs, allowing for
improved performance compared to executing the Python
function directly in a loop.
The print(result) statement simply prints the result obtained
from invoking the TensorFlow graph function with an input
value of 2.