TensorFlow
benchmarking
graph performance
machine learning
optimization

What is the proper way to benchmark part of tensorflow graph?

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

Introduction

Benchmarking part of a TensorFlow graph is only useful if you isolate the work you care about, warm it up first, and force execution to complete before reading the timer. Timing a single call without synchronization often measures tracing, compilation, or queueing overhead instead of the actual operation. The proper method is to wrap the target computation, run several warm-up iterations, then average repeated measured runs.

Isolate the Part You Want to Measure

Start by putting the specific subgraph or operation into its own function. In TensorFlow 2, tf.function is usually the closest equivalent to graph execution.

python
1import tensorflow as tf
2
3@tf.function
4def target_op(x):
5    x = tf.matmul(x, x)
6    x = tf.nn.relu(x)
7    return x

This keeps the benchmark focused. If you time a huge training step when you only care about one matrix multiply block, the result is not actionable.

Warm Up Before Measuring

The first few runs are often not representative because TensorFlow may trace the function, build graphs, or trigger one-time setup costs.

python
1import tensorflow as tf
2
3x = tf.random.normal((1024, 1024))
4
5for _ in range(5):
6    _ = target_op(x)

Without warm-up, the first timing can significantly overstate steady-state runtime.

Force Execution to Finish

TensorFlow operations, especially on accelerators, may execute asynchronously. A timer around the function call can stop before the real work completes unless you force synchronization.

One simple method is to materialize the result:

python
result = target_op(x)
_ = result.numpy()

Calling .numpy() forces completion in eager execution and makes the timing more trustworthy.

Measure Repeated Runs

Use repeated runs and average them rather than trusting one sample.

python
1import time
2import tensorflow as tf
3
4x = tf.random.normal((1024, 1024))
5
6for _ in range(5):
7    _ = target_op(x).numpy()
8
9times = []
10for _ in range(20):
11    start = time.perf_counter()
12    _ = target_op(x).numpy()
13    end = time.perf_counter()
14    times.append(end - start)
15
16print("avg ms:", 1000 * sum(times) / len(times))
17print("min ms:", 1000 * min(times))

Average and minimum together often tell a better story than a single number.

Benchmark the Same Input Shape You Use in Reality

TensorFlow performance depends heavily on tensor shapes, dtypes, and device placement. A benchmark using tiny toy inputs may say nothing useful about the production workload.

Make sure the benchmark matches:

  • input shape
  • batch size
  • dtype
  • CPU or GPU device

If you plan to optimize a real inference path, benchmark the real inference shape.

Use TensorBoard Profiler for Deeper Analysis

Wall-clock timing tells you how long something takes. It does not tell you why. For deeper analysis, use TensorFlow profiling.

python
1import tensorflow as tf
2
3tf.profiler.experimental.start("logdir")
4for _ in range(10):
5    _ = target_op(x).numpy()
6tf.profiler.experimental.stop()

You can then inspect the trace in TensorBoard to see kernel launches, device utilization, and operator-level timing.

This is the right next step once a simple benchmark shows there is a real performance issue.

Legacy TensorFlow 1 Graphs

If you are working with TensorFlow 1 style sessions, the same principles apply:

  • isolate the tensor or op to run
  • warm up the session
  • call sess.run repeatedly
  • time only the steady-state execution

Example structure:

python
1import time
2import tensorflow as tf
3
4tf.compat.v1.disable_eager_execution()
5
6graph = tf.Graph()
7with graph.as_default():
8    x = tf.compat.v1.placeholder(tf.float32, shape=(1024, 1024))
9    y = tf.matmul(x, x)
10
11with tf.compat.v1.Session(graph=graph) as sess:
12    data = tf.random.normal((1024, 1024)).numpy()
13
14    for _ in range(5):
15        sess.run(y, feed_dict={x: data})
16
17    start = time.perf_counter()
18    for _ in range(20):
19        sess.run(y, feed_dict={x: data})
20    end = time.perf_counter()
21
22    print("avg ms:", 1000 * (end - start) / 20)

The methodology stays the same even though the API style is older.

Common Pitfalls

The biggest mistake is timing the first call and treating it as steady-state performance. Graph tracing and one-time setup can dominate that result.

Another issue is forgetting asynchronous execution on accelerators. If you do not force completion, you may benchmark submission time rather than execution time.

Developers also often benchmark unrealistic input shapes, then optimize the wrong thing because the benchmark never represented the real workload.

Summary

  • Isolate the exact TensorFlow subgraph or function you want to measure.
  • Warm it up before collecting timings.
  • Force execution to complete before stopping the timer.
  • Measure repeated runs and use representative input shapes.
  • Use TensorBoard profiling when wall-clock timing alone is not enough.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the 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

All Rights Reserved.