Compare static graph using @tf.function VS dynamic graph.
AutoGraph helps you write complicated graph code using normal Python. Behind the scenes, AutoGraph automatically transforms your code into the equivalent TensorFlow graph code.
Let's take a look at TensorFlow graphs and how they work.
ReLU_Layer = tf.keras.layers.Dense(100, input_shape=(784,), activation=tf.nn.relu)
Logit_Layer = tf.keras.layers.Dense(10, input_shape=(100,))
# X and y are labels and inputs
TensorFlow 1.0: Operations are added as nodes to the computational graph and are not actually executed until we call session.run(), much like defining a function that doesn't run until it is called.
SGD_Trainer = tf.train.GradientDescentOptimizer(1e-2)
inputs = tf.placeholder(tf.float32, shape=[None, 784])
labels = tf.placeholder(tf.int16, shape=[None, 10])
hidden = ReLU_Layer(inputs)
logits = Logit_Layer(hidden)
entropy = tf.nn.softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
loss = tf.reduce_mean(entropy)
train_step = SGD_Trainer.minimize(loss,
var_list=ReLU_Layer.weights+Logit_Layer.weights)
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
for step in range(1000):
sess.run(train_step, feed_dict={inputs:X, labels:y})
TensorFlow 2.0: Operations are executed directly and the computational graph is built on-the-fly. However, we can still write functions and pre-compile computational graphs from them like in TF 1.0 using the @tf.function decorator, allowing for faster execution.
SGD_Trainer = tf.optimizers.SGD(1e-2)
@tf.function
def loss_fn(inputs=X, labels=y):
hidden = ReLU_Layer(inputs)
logits = Logit_Layer(hidden)
entropy = tf.nn.softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
return tf.reduce_mean(entropy)
for step in range(1000):
SGD_Trainer.minimize(loss_fn,
var_list=ReLU_Layer.weights+Logit_Layer.weights)
python main.py
and you will see some computation cost between static graph and dynamic graph.