Speed Test: Eager vs Graph Mode Benchmark in TensorFlow

The Story — A Fair Race

Everyone says "graphs are faster". In this lesson you will measure it yourself, and find out when it is true — and when it is not. All numbers below were measured on our small 2-core CPU test machine. Your numbers will be different, but the pattern will be the same.

How to Time TensorFlow Code Fairly

import tensorflow as tf
import timeit

def bench(fn, *args, number=100):
    fn(*args)                                              # 1. warm-up (traces if needed)
    run = lambda: fn(*args).numpy()                        # 3. wait for the result
    best = min(timeit.repeat(run, number=number, repeat=3))   # 2 + 4
    return best / number * 1e6                             # microseconds per call

print("helper ready")

Test 1 — Many Small Ops

A loop of 50 steps, 3 tiny ops each, on a 10 × 10 tensor. This is where eager mode suffers most: 150 trips between Python and the runtime.

def small_ops(x):
    for _ in range(50):
        x = tf.nn.relu(x * 0.99 + 0.01)
    return x

x = tf.random.normal([10, 10])
graph_small = tf.function(small_ops)

eager_us = bench(small_ops, x)
graph_us = bench(graph_small, x)
print(f"eager : {eager_us:8.1f} µs per call")
print(f"graph : {graph_us:8.1f} µs per call")
print(f"speed-up: {eager_us / graph_us:.1f}x")

Test 2 — One Big Op

Now the opposite: a single 1000 × 1000 matrix multiply. The maths itself takes most of the time, so there is little Python overhead to remove.

def big_matmul(a):
    return tf.matmul(a, a)

a = tf.random.normal([1000, 1000])
graph_big = tf.function(big_matmul)

eager_us = bench(big_matmul, a, number=10)
graph_us = bench(graph_big, a, number=10)
print(f"eager : {eager_us / 1000:6.2f} ms per call")
print(f"graph : {graph_us / 1000:6.2f} ms per call")
print(f"speed-up: {eager_us / graph_us:.2f}x")

Test 3 — A Real Training Step

This is the case that matters most. A training step has a forward pass, a loss, a backward pass and an optimizer update — dozens of small ops.

tf.random.set_seed(0)
model = tf.keras.Sequential([
    tf.keras.Input(shape=(20,)),
    tf.keras.layers.Dense(64, activation="relu"),
    tf.keras.layers.Dense(1),
])
opt = tf.keras.optimizers.SGD(0.01)
X, Y = tf.random.normal([64, 20]), tf.random.normal([64, 1])

def train_step(X, Y):
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((model(X) - Y) ** 2)
    grads = tape.gradient(loss, model.trainable_variables)
    opt.apply_gradients(zip(grads, model.trainable_variables))
    return loss

graph_step = tf.function(train_step)

eager_us = bench(train_step, X, Y, number=50)
graph_us = bench(graph_step, X, Y, number=50)
print(f"eager step : {eager_us / 1000:6.2f} ms")
print(f"graph step : {graph_us / 1000:6.2f} ms")
print(f"speed-up   : {eager_us / graph_us:.1f}x   (Keras fit() does this for you)")

Test 4 — The Cost of Tracing

The first call is slow because it traces. That cost is paid once per signature — and again on every retrace.

import time

fresh = tf.function(train_step)
t0 = time.perf_counter(); fresh(X, Y).numpy(); first = time.perf_counter() - t0
t0 = time.perf_counter(); fresh(X, Y).numpy(); second = time.perf_counter() - t0

print(f"1st call (trace + run): {first * 1000:7.1f} ms")
print(f"2nd call (run only)   : {second * 1000:7.2f} ms")
print(f"tracing costs about {first / second:.0f} normal calls")

Bonus: XLA Compilation

jit_compile=True asks the XLA compiler to fuse ops into bigger kernels. It can help more, but test it — it does not always win.

xla_small = tf.function(small_ops, jit_compile=True)
print(f"graph     : {bench(graph_small, x):6.1f} µs")
print(f"graph+XLA : {bench(xla_small, x):6.1f} µs")

Where the Time Goes

Try It — Speed Simulator

When Is tf.function Worth It?

Golden Rules