Practice Now
TensorFlow
TensorFlow 2.0
eager execution
machine learning
deep learning

Using Tensorflow 2.0 and eager execution without Keras

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Deep LearningTensorFlow & KerasMachine Learning

Introduction

TensorFlow 2.0 enables eager execution by default, allowing operations to execute immediately like normal Python code. While Keras is the recommended high-level API, TensorFlow 2.0 provides powerful low-level APIs for building models without Keras — using tf.Variable, tf.GradientTape, and tf.Module. This approach gives full control over the training loop, custom gradient computation, and model architecture, which is essential for research, custom layers, and non-standard training procedures.

Basic Tensor Operations in Eager Mode

python
1import tensorflow as tf
2
3# Operations execute immediately — no sessions needed
4a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
5b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
6
7# Matrix operations
8c = tf.matmul(a, b)
9print(c)  # tf.Tensor([[19. 22.] [43. 50.]], shape=(2, 2), dtype=float32)
10
11# Convert to numpy
12print(c.numpy())  # [[19. 22.] [43. 50.]]
13
14# Element-wise operations
15print(tf.reduce_mean(a))  # 2.5
16print(tf.reduce_sum(b, axis=1))  # [11. 15.]

In eager mode, every tf operation returns a concrete value immediately. You can inspect tensors, use Python control flow (if, for), and debug with standard Python tools.

Variables and GradientTape

python
1import tensorflow as tf
2
3# Variables hold mutable state
4w = tf.Variable(tf.random.normal([3, 1]), name="weights")
5b = tf.Variable(tf.zeros([1]), name="bias")
6
7# GradientTape records operations for automatic differentiation
8x = tf.constant([[1.0, 2.0, 3.0]])
9y_true = tf.constant([[7.0]])
10
11with tf.GradientTape() as tape:
12    y_pred = tf.matmul(x, w) + b
13    loss = tf.reduce_mean((y_true - y_pred) ** 2)
14
15# Compute gradients
16gradients = tape.gradient(loss, [w, b])
17print(f"dL/dw: {gradients[0].numpy().flatten()}")
18print(f"dL/db: {gradients[1].numpy()}")

tf.GradientTape records all operations involving tf.Variable inside its context. Calling tape.gradient(loss, variables) computes partial derivatives using reverse-mode autodiff.

Building a Model with tf.Module

python
1import tensorflow as tf
2
3class LinearModel(tf.Module):
4    def __init__(self, input_dim, output_dim, name="linear"):
5        super().__init__(name=name)
6        self.w = tf.Variable(
7            tf.random.normal([input_dim, output_dim]), name="weights"
8        )
9        self.b = tf.Variable(tf.zeros([output_dim]), name="bias")
10
11    def __call__(self, x):
12        return tf.matmul(x, self.w) + self.b
13
14class MLP(tf.Module):
15    def __init__(self, name="mlp"):
16        super().__init__(name=name)
17        self.layer1 = LinearModel(784, 128, name="hidden")
18        self.layer2 = LinearModel(128, 10, name="output")
19
20    def __call__(self, x):
21        x = tf.nn.relu(self.layer1(x))
22        return self.layer2(x)
23
24model = MLP()
25# Variables are tracked automatically by tf.Module
26print(f"Trainable variables: {len(model.trainable_variables)}")
27# 4 (w1, b1, w2, b2)

tf.Module automatically tracks tf.Variable attributes in submodules. It provides .trainable_variables without manual bookkeeping.

Custom Training Loop

python
1import tensorflow as tf
2import numpy as np
3
4# Generate synthetic data
5np.random.seed(42)
6X_train = np.random.randn(1000, 784).astype(np.float32)
7y_train = np.random.randint(0, 10, 1000)
8
9# Create model and optimizer
10model = MLP()
11optimizer = tf.optimizers.Adam(learning_rate=0.001)
12loss_fn = tf.nn.sparse_softmax_cross_entropy_with_logits
13
14# Training loop
15batch_size = 32
16epochs = 5
17
18for epoch in range(epochs):
19    epoch_loss = 0.0
20    num_batches = 0
21
22    for i in range(0, len(X_train), batch_size):
23        x_batch = tf.constant(X_train[i:i+batch_size])
24        y_batch = tf.constant(y_train[i:i+batch_size])
25
26        with tf.GradientTape() as tape:
27            logits = model(x_batch)
28            loss = tf.reduce_mean(loss_fn(labels=y_batch, logits=logits))
29
30        gradients = tape.gradient(loss, model.trainable_variables)
31        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
32
33        epoch_loss += loss.numpy()
34        num_batches += 1
35
36    print(f"Epoch {epoch+1}: loss = {epoch_loss / num_batches:.4f}")

The custom training loop gives full control over batching, gradient computation, gradient clipping, learning rate scheduling, and logging — all in plain Python.

Using tf.function for Performance

python
1import tensorflow as tf
2
3model = MLP()
4optimizer = tf.optimizers.Adam(0.001)
5
6@tf.function  # compile to a TF graph for speed
7def train_step(x, y):
8    with tf.GradientTape() as tape:
9        logits = model(x)
10        loss = tf.reduce_mean(
11            tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=logits)
12        )
13    gradients = tape.gradient(loss, model.trainable_variables)
14    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
15    return loss
16
17# First call traces the function — subsequent calls use the compiled graph
18x_batch = tf.random.normal([32, 784])
19y_batch = tf.random.uniform([32], maxval=10, dtype=tf.int32)
20
21loss = train_step(x_batch, y_batch)
22print(f"Loss: {loss.numpy():.4f}")

@tf.function traces the Python function once and compiles it into an optimized TensorFlow graph. Subsequent calls execute the graph directly, bypassing Python overhead. This provides significant speedups for training and inference.

Saving and Loading Without Keras

python
1import tensorflow as tf
2
3model = MLP()
4# Build the model with a dummy forward pass
5model(tf.zeros([1, 784]))
6
7# Save with tf.train.Checkpoint
8checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
9checkpoint.save("/tmp/my_model/ckpt")
10
11# Restore
12new_model = MLP()
13new_optimizer = tf.optimizers.Adam(0.001)
14new_checkpoint = tf.train.Checkpoint(model=new_model, optimizer=new_optimizer)
15new_checkpoint.restore(tf.train.latest_checkpoint("/tmp/my_model"))
16
17# Verify
18new_model(tf.zeros([1, 784]))  # works after restore
19
20# Save as SavedModel (for serving)
21tf.saved_model.save(model, "/tmp/saved_model")
22loaded = tf.saved_model.load("/tmp/saved_model")

Common Pitfalls

  • Forgetting tf.GradientTape context for gradient computation: Operations outside the with tf.GradientTape() block are not recorded. If the forward pass happens outside the tape, tape.gradient() returns None for all variables. Ensure the entire forward pass and loss computation are inside the tape context.
  • Using Python scalars instead of tensors in @tf.function: @tf.function traces the function once per unique input signature. Passing Python integers or floats causes retracing on every call. Convert to tf.constant or use tensor arguments to avoid performance degradation from repeated tracing.
  • Not calling the model before saving: tf.Module variables are created lazily during the first forward pass. Saving a model that has never been called saves an empty checkpoint. Always run a dummy forward pass (model(tf.zeros([1, input_dim]))) before saving.
  • Modifying variables outside tf.GradientTape: Direct variable assignment (w.assign(new_value)) outside the tape is not tracked for gradient computation. Use the tape to record operations that should contribute to gradients, and use optimizer.apply_gradients() for parameter updates.
  • Assuming eager mode is always slower than graph mode: While eager execution has Python overhead per operation, @tf.function eliminates most of it by compiling to a graph. For training loops, wrapping the training step in @tf.function provides near-identical performance to TF1 graph mode.

Summary

  • TF2 eager mode executes operations immediately without sessions or graphs
  • Use tf.Variable for model parameters and tf.GradientTape for automatic differentiation
  • Subclass tf.Module to organize variables and build modular models
  • Write custom training loops with GradientTape + optimizer.apply_gradients()
  • Use @tf.function to compile training steps into optimized graphs for production speed
  • Save models with tf.train.Checkpoint (for training) or tf.saved_model.save (for serving)
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.