TensorFlow architecture (eager vs graph execution) - Experiment Comparison

Experiment - TensorFlow architecture (eager vs graph execution)

Problem:You have a simple neural network model built with TensorFlow using eager execution. The model trains correctly but runs slower than expected on larger datasets.

Current Metrics:Training accuracy: 85%, Validation accuracy: 83%, Training time per epoch: 12 seconds

Issue:The model runs slower because eager execution evaluates operations immediately, which is easier to debug but less optimized for speed.

Your Task

Improve the training speed by switching from eager execution to graph execution while maintaining similar accuracy (validation accuracy >= 80%).

Keep the model architecture the same.

Do not change the dataset or batch size.

Use TensorFlow 2.x features only.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers, models

# Load dataset
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# Define model architecture
class SimpleModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.flatten = layers.Flatten()
        self.dense1 = layers.Dense(128, activation='relu')
        self.dense2 = layers.Dense(10)

    @tf.function  # Convert call method to graph execution
    def call(self, x):
        x = self.flatten(x)
        x = self.dense1(x)
        return self.dense2(x)

model = SimpleModel()

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()

train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()

@tf.function  # Graph execution for training step
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x)
        loss = loss_fn(y, logits)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_acc_metric.update_state(y, logits)
    return loss

# Training loop
batch_size = 64
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)

for epoch in range(3):
    train_acc_metric.reset_state()
    for batch_x, batch_y in train_dataset:
        loss = train_step(batch_x, batch_y)
    train_acc = train_acc_metric.result()
    print(f"Epoch {epoch + 1}, Loss: {loss:.4f}, Training Accuracy: {train_acc:.4f}")

# Evaluate on test data
logits = model(x_test)
predictions = tf.argmax(logits, axis=1, output_type=tf.int32)
accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, y_test), tf.float32))
print(f"Validation Accuracy: {accuracy:.4f}")

Added @tf.function decorator to the model's call method to enable graph execution.

Wrapped the training step function with @tf.function to compile it into a graph.

Kept the model architecture and dataset unchanged to isolate the effect of execution mode.

Results Interpretation

Before: Training accuracy 85%, Validation accuracy 83%, Training time 12s per epoch.

After: Training accuracy 84%, Validation accuracy 82%, Training time 5s per epoch.

Switching from eager to graph execution in TensorFlow speeds up training significantly with minimal impact on accuracy. Graph execution compiles operations into an optimized graph, making the model run faster especially on larger datasets.

Bonus Experiment

Try adding tf.function to only the training step but keep the model call method in eager mode. Compare speed and accuracy.

💡 Hint

This will show how partial graph execution affects performance and debugging ease.

Practice

1. What is the main difference between eager execution and graph execution in TensorFlow?

easy

A. Eager execution requires a GPU, graph execution runs only on CPU.

B. Eager execution uses less memory than graph execution in all cases.

C. Graph execution is only for training, eager execution is only for inference.

D. Eager execution runs operations immediately, while graph execution builds a computation plan first.

Solution

1. Step 1: Understand eager execution behavior

   Eager execution runs TensorFlow operations immediately as they are called, making it easy to debug and understand.

2. Step 2: Understand graph execution behavior

   Graph execution builds a computation graph first, then runs it for better performance and optimization.

3. Final Answer:

   Eager execution runs operations immediately, while graph execution builds a computation plan first. -> Option D

4. Quick Check:

   Eager vs Graph = Immediate vs Plan [OK]

Hint: Eager means now, graph means plan first [OK]

Common Mistakes:

• Thinking graph execution runs immediately
• Confusing hardware requirements
• Assuming eager is only for inference

2. Which of the following is the correct way to convert a Python function to a TensorFlow graph function?

easy

A. Use @tf.function decorator above the function definition.

B. Call tf.convert_to_graph(function) before running it.

C. Wrap the function inside tf.Graph() and call it.

D. Set tf.enable_graph_mode(True) before defining the function.

Solution

1. Step 1: Recall TensorFlow's method to switch execution modes

   TensorFlow uses the @tf.function decorator to convert a Python function into a graph function.

2. Step 2: Evaluate other options for correctness

   tf.convert_to_graph and tf.enable_graph_mode do not exist; wrapping in tf.Graph() is not the standard way.

3. Final Answer:

   Use @tf.function decorator above the function definition. -> Option A

4. Quick Check:

   @tf.function converts to graph [OK]

Hint: Remember @tf.function for graph conversion [OK]

Common Mistakes:

• Using non-existent TensorFlow functions
• Trying to enable graph mode globally
• Confusing tf.Graph() usage

3. Consider the following code snippet:

import tensorflow as tf

@tf.function
def add(a, b):
    print('Running add')
    return a + b

result1 = add(1, 2)
result2 = add(3, 4)

What will be printed when this code runs?

medium

A. Running add
Running add

B. Running add

C. No output printed

D. Error due to print inside @tf.function

Solution

1. Step 1: Understand print behavior inside @tf.function

   When a function is decorated with @tf.function, it runs as a graph. Python print runs only once during graph tracing, not on every call.

2. Step 2: Analyze the calls to add()

   The first call triggers tracing and prints 'Running add'. The second call uses the compiled graph and does not print again.

3. Final Answer:

   Running add -> Option B

4. Quick Check:

   Print runs once during tracing [OK]

Hint: Print inside @tf.function runs once [OK]

Common Mistakes:

• Expecting print every call
• Thinking print is disabled
• Assuming error from print usage

4. You wrote this code:

import tensorflow as tf

def multiply(a, b):
    return a * b

@tf.function
def call_multiply(x, y):
    return multiply(x, y)

print(call_multiply(2, 3))

But the output is a Tensor object, not a number. How can you fix it to print the actual number?

medium

A. Wrap multiply inside tf.function as well

B. Remove @tf.function decorator from call_multiply

C. Add .numpy() to the print call: print(call_multiply(2, 3).numpy())

D. Change multiply to use tf.multiply instead of * operator

Solution

1. Step 1: Understand output type of @tf.function

   Functions decorated with @tf.function return TensorFlow tensors, not plain Python numbers.

2. Step 2: Convert tensor to number for printing

   Use the .numpy() method on the tensor to get the actual number value for printing.

3. Final Answer:

   Add .numpy() to the print call: print(call_multiply(2, 3).numpy()) -> Option C

4. Quick Check:

   Tensor to number: use .numpy() [OK]

Hint: Use .numpy() to get number from tensor [OK]

Common Mistakes:

• Expecting tensor to print as number
• Removing @tf.function unnecessarily
• Changing multiply without need

5. You want to speed up a TensorFlow model training loop by switching from eager to graph execution. Which approach correctly applies this change while keeping eager mode for debugging?

hard

A. Decorate the training step function with @tf.function and run training normally.

B. Set tf.config.experimental_run_functions_eagerly(True) before training.

C. Rewrite the entire model using tf.Graph() and tf.Session().

D. Disable eager execution globally using tf.compat.v1.disable_eager_execution().

Solution

1. Step 1: Identify how to switch to graph execution selectively

   Using @tf.function on the training step compiles it to a graph, speeding execution while keeping eager mode elsewhere.

2. Step 2: Evaluate other options for drawbacks

   Setting experimental_run_functions_eagerly(True) forces eager mode (slower). Rewriting with tf.Graph() and tf.Session() is outdated. Disabling eager globally removes debugging ease.

3. Final Answer:

   Decorate the training step function with @tf.function and run training normally. -> Option A

4. Quick Check:

   @tf.function speeds training, keeps eager debugging [OK]

Hint: Use @tf.function on training step for speed [OK]

Common Mistakes:

• Forcing eager mode instead of graph
• Using old TensorFlow 1.x APIs
• Disabling eager globally losing debugging