TensorFlow Eager Execution

Introduction

Eager Execution is a programming environment in TensorFlow that evaluates operations immediately, without building computational graphs. This represents a significant shift from TensorFlow's original design, which required developers to first build a computational graph and then execute it within a session.

Introduced in TensorFlow 1.x and becoming the default mode in TensorFlow 2.x, Eager Execution makes TensorFlow more intuitive, especially for beginners. It allows you to write code in a more imperative, Python-like manner and see results immediately, making debugging and prototyping much easier.

In this tutorial, you'll learn:

• What Eager Execution is and how it differs from graph-based execution
• How to enable and use Eager Execution
• Benefits and potential drawbacks
• Practical examples demonstrating its use in real scenarios

Understanding Eager Execution

Graph-based vs Eager Execution

Before Eager Execution, TensorFlow worked exclusively in a graph-based mode:

1. Graph-based execution (Traditional TensorFlow 1.x approach):

   • Define operations to build a computational graph
   • Create a session
   • Run the graph within the session to get results

2. Eager Execution (Default in TensorFlow 2.x):

   • Operations are evaluated immediately
   • Values are returned directly
   • No need for session creation or graph building

Let's see how the same operation looks in both modes:

python

# Graph-based execution (TensorFlow 1.x style)
import tensorflow as tf

# Create a graph
a = tf.constant(5.0)
b = tf.constant(6.0)
c = a * b

# Need a session to evaluate
sess = tf.Session()
result = sess.run(c)
print(result)  # 30.0
sess.close()

# Eager execution (TensorFlow 2.x style)
import tensorflow as tf
tf.compat.v1.enable_eager_execution()  # Not needed in TF 2.x

a = tf.constant(5.0)
b = tf.constant(6.0)
c = a * b

print(c)  # tf.Tensor(30.0, shape=(), dtype=float32)

Enabling Eager Execution

In TensorFlow 2.x

In TensorFlow 2.x, Eager Execution is enabled by default. You don't need to do anything special:

python

import tensorflow as tf

# Check if eager execution is enabled
print(f"Eager execution enabled: {tf.executing_eagerly()}")

# Simple operations execute immediately
x = tf.constant([[1, 2], [3, 4]])
print(x)
print(x.numpy())  # Convert to NumPy array

Output:

Eager execution enabled: True
tf.Tensor(
[[1 2]
 [3 4]], shape=(2, 2), dtype=int32)
[[1 2]
 [3 4]]

In TensorFlow 1.x

If you're using TensorFlow 1.x, you need to explicitly enable Eager Execution:

python

import tensorflow as tf
tf.compat.v1.enable_eager_execution()

# Verify it's enabled
print(f"Eager execution enabled: {tf.executing_eagerly()}")

Important: Eager Execution must be enabled at the beginning of your program, before any other TensorFlow operations.

Key Features of Eager Execution

1. Dynamic Control Flow

With Eager Execution, you can use Python control flow (if statements, loops, etc.) directly with TensorFlow operations:

python

import tensorflow as tf
import numpy as np

def simple_nn(x, y):
    w = tf.Variable([[1.0]])
    with tf.GradientTape() as tape:
        prediction = tf.matmul(x, w)
        loss = tf.reduce_mean(tf.square(prediction - y))
    
    # Get gradients
    gradients = tape.gradient(loss, [w])
    
    # Print information during training
    if loss < 0.5:
        print(f"Loss is getting smaller: {loss.numpy()}")
    
    return loss, gradients

# Sample data
x_train = np.array([[1.0], [2.0], [3.0], [4.0]])
y_train = np.array([[2.0], [4.0], [6.0], [8.0]])

# Convert to tensors
x = tf.constant(x_train)
y = tf.constant(y_train)

loss, gradients = simple_nn(x, y)
print(f"Loss: {loss.numpy()}")
print(f"Gradients: {gradients[0].numpy()}")

Output:

Loss: 0.0
Gradients: [[-0.]]

2. Immediate Feedback and Debugging

One of the biggest advantages of Eager Execution is the ability to inspect variables and debug immediately:

python

import tensorflow as tf

# Create some tensors
x = tf.random.normal([3, 3])
print("Original tensor:")
print(x.numpy())

# Apply operations
y = tf.nn.relu(x)  # Apply ReLU activation
print("\nAfter ReLU (negative values become 0):")
print(y.numpy())

# Use Python's debugging tools
import pdb

def debug_function():
    a = tf.constant([1, 2])
    b = tf.constant([3, 4])
    c = a + b
    # Uncomment the next line to use the debugger
    # pdb.set_trace()
    d = c * 2
    return d

result = debug_function()
print("\nDebuggable result:")
print(result.numpy())

The output will show the tensor values at each step, making it much easier to understand what's happening.

3. Custom Gradients and Automatic Differentiation

Eager Execution works seamlessly with TensorFlow's automatic differentiation:

python

import tensorflow as tf

# Define a simple function
def f(x):
    return tf.square(x)

# Use GradientTape to calculate derivatives
x = tf.Variable(3.0)
with tf.GradientTape() as tape:
    y = f(x)
    
# dy/dx = d(x²)/dx = 2x
dy_dx = tape.gradient(y, x)
print(f"Value of x: {x.numpy()}")
print(f"Value of f(x) = x²: {y.numpy()}")
print(f"Derivative of f(x) at x=3: {dy_dx.numpy()}")  # Should be 2*3 = 6

Output:

Value of x: 3.0
Value of f(x) = x²: 9.0
Derivative of f(x) at x=3: 6.0

Practical Applications

Example 1: Linear Regression with Eager Execution

Let's implement a simple linear regression model using Eager Execution:

python

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# Generate synthetic data
np.random.seed(42)
x = np.random.rand(100, 1)
y = 2 * x + 1 + 0.1 * np.random.randn(100, 1)  # y = 2x + 1 + noise

# Convert to tensors
x_tensor = tf.convert_to_tensor(x, dtype=tf.float32)
y_tensor = tf.convert_to_tensor(y, dtype=tf.float32)

# Initialize parameters
W = tf.Variable(np.random.randn(), name='weight')
b = tf.Variable(np.random.randn(), name='bias')

# Define the model
def linear_model(x):
    return W * x + b

# Define loss function
def loss_fn(y_pred, y_true):
    return tf.reduce_mean(tf.square(y_pred - y_true))

# Training parameters
learning_rate = 0.1
epochs = 100

# Optimizer
optimizer = tf.keras.optimizers.SGD(learning_rate)

# Training loop
losses = []
for epoch in range(epochs):
    with tf.GradientTape() as tape:
        predictions = linear_model(x_tensor)
        loss = loss_fn(predictions, y_tensor)
    
    gradients = tape.gradient(loss, [W, b])
    optimizer.apply_gradients(zip(gradients, [W, b]))
    
    losses.append(loss.numpy())
    
    if epoch % 20 == 0:
        print(f"Epoch {epoch}, Loss: {loss.numpy():.4f}, W: {W.numpy():.4f}, b: {b.numpy():.4f}")

print(f"Final parameters - Weight: {W.numpy():.4f}, Bias: {b.numpy():.4f}")

Output:

Epoch 0, Loss: 1.8461, W: 0.8018, b: 0.3877
Epoch 20, Loss: 0.0345, W: 1.8076, b: 0.9303
Epoch 40, Loss: 0.0152, W: 1.9178, b: 0.9940
Epoch 60, Loss: 0.0128, W: 1.9601, b: 1.0182
Epoch 80, Loss: 0.0121, W: 1.9785, b: 1.0284
Final parameters - Weight: 1.9868, Bias: 1.0329

The model learns parameters close to the true values (W=2, b=1).

Example 2: Custom Training Loop for Image Classification

Let's implement a simple custom training loop for image classification on the MNIST dataset:

python

import tensorflow as tf

# Load and preprocess the MNIST 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  # Normalize pixel values

# Reshape and cast
x_train = x_train.reshape(-1, 28*28).astype('float32')
x_test = x_test.reshape(-1, 28*28).astype('float32')

# Convert labels to one-hot encoding
y_train = tf.one_hot(y_train, 10)
y_test = tf.one_hot(y_test, 10)

# Create a simple model
class MNISTModel(tf.keras.Model):
    def __init__(self):
        super(MNISTModel, self).__init__()
        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10)
    
    def call(self, x):
        x = self.dense1(x)
        return self.dense2(x)

model = MNISTModel()
loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()

# Metrics
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy')

# Training step function
@tf.function  # This decorator compiles the function into a graph for faster execution
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_object(labels, predictions)
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    train_loss(loss)
    train_accuracy(labels, predictions)

# Training loop
batch_size = 32
epochs = 5

for epoch in range(epochs):
    # Reset the metrics
    train_loss.reset_states()
    train_accuracy.reset_states()
    
    # Create batches
    dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
    
    for images, labels in dataset:
        train_step(images, labels)
    
    print(f'Epoch {epoch+1}, '
          f'Loss: {train_loss.result():.4f}, '
          f'Accuracy: {train_accuracy.result()*100:.2f}%')

# Test the model
test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.CategoricalAccuracy(name='test_accuracy')

test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)
for test_images, test_labels in test_dataset:
    predictions = model(test_images)
    t_loss = loss_object(test_labels, predictions)
    
    test_loss(t_loss)
    test_accuracy(test_labels, predictions)

print(f'Test Loss: {test_loss.result():.4f}, Test Accuracy: {test_accuracy.result()*100:.2f}%')

This example demonstrates a custom training loop with Eager Execution, including:

• Creating a custom model class
• Implementing a training step function
• Using metrics to track progress
• The @tf.function decorator for compiling parts of the code into graphs for faster execution

Benefits and Considerations

Advantages of Eager Execution

1. Easier Debugging: Inspect values as they're computed
2. Python Control Flow: Use native Python control flow statements with TensorFlow operations
3. Natural Programming Style: Write more intuitive code similar to normal Python
4. Immediate Error Feedback: Get errors at operation execution time instead of graph construction time
5. Simpler Prototyping: Quickly test ideas without building complex graphs

Considerations

1. Performance: Eager Execution can be slower than graph execution for large models
2. Memory Usage: May use more memory due to immediate value computation
3. Production Deployment: For production, consider using @tf.function to compile critical parts to graphs

The @tf.function Decorator

TensorFlow 2.x offers the best of both worlds with the @tf.function decorator, which allows you to convert eager execution code into graph execution for better performance:

python

import tensorflow as tf
import time

# Regular Python function (eager execution)
def eager_function(x):
    return tf.reduce_sum(tf.square(x))

# Graph-compiled function
@tf.function
def graph_function(x):
    return tf.reduce_sum(tf.square(x))

# Compare performance
x = tf.random.normal([1000, 1000])

# Warmup
_ = eager_function(x)
_ = graph_function(x)

# Test eager execution
start_time = time.time()
for _ in range(100):
    _ = eager_function(x)
eager_time = time.time() - start_time

# Test graph execution
start_time = time.time()
for _ in range(100):
    _ = graph_function(x)
graph_time = time.time() - start_time

print(f"Eager execution time: {eager_time:.4f} seconds")
print(f"Graph execution time: {graph_time:.4f} seconds")
print(f"Speedup: {eager_time / graph_time:.2f}x")

The graph-compiled version will typically run significantly faster for repeated operations.

Summary

In this tutorial, you've learned about TensorFlow's Eager Execution mode:

• Eager Execution allows for immediate evaluation of operations, making TensorFlow code more intuitive and easier to debug
• It's the default mode in TensorFlow 2.x, making the framework more accessible for beginners
• Key features include Python control flow integration, immediate feedback, and seamless automatic differentiation
• For performance-critical applications, the @tf.function decorator allows you to compile parts of your code into graphs
• Practical examples showed how to implement linear regression and custom training loops using this paradigm

With Eager Execution, TensorFlow becomes more accessible to those familiar with imperative programming while maintaining the performance benefits of graph execution when needed.

Additional Resources

• TensorFlow Official Guide on Eager Execution
• Introduction to Graphs and tf.function
• Better Performance with tf.function

Exercises

1. Create a simple neural network for the XOR problem using Eager Execution
2. Implement a custom gradient in a neural network layer
3. Convert an existing graph-based TensorFlow 1.x model to use Eager Execution
4. Profile the performance of a model running with Eager Execution vs. using @tf.function
5. Build a custom training loop for a convolutional neural network on the CIFAR-10 dataset

By working through these exercises, you'll develop a deeper understanding of Eager Execution and how to leverage it effectively in your TensorFlow projects.