TensorFlow GPU vs CPU

A GPU does not make a model faster; it makes large parallel tensor ops faster and leaves everything else unchanged. The difference matters when you know which parts of a training step are matrix multiplies and which are Python, I/O, or tiny ops. Measuring is the only reliable way to decide.

The idea

• A CPU has tens of cores optimized for latency; a GPU has thousands optimized for throughput. Matrix multiplication of (b, m, k) @ (k, n) maps perfectly onto that throughput.
• The speedup is real only when tensors are large enough to saturate the device. A (32, 28, 28, 1) batch is small; a (512, 1024) @ (1024, 1024) GEMM is not.
• Device placement matters. with tf.device("/GPU:0") pins ops; unplaced math may silently run on the CPU.
• The pipeline can become the bottleneck: if tf.data cannot feed batches fast enough the GPU idles, and the comparison measures your loader, not the chip.

Worked example

Time the same (512, 1024) @ (1024, 1024) matmul 200 times.

• Warm up first; the initial kernel launch pays setup cost that would skew the average.
• Check devices with tf.config.list_physical_devices("GPU").
• Wrap the loop in time.perf_counter() and call .numpy() once at the end to force synchronization — without it you measure kernel launch, not execution.
• Report throughput as 2 * m * n * k * iters / elapsed, the standard GEMM FLOP count.

For this shape the GEMM has 2 × 512 × 1024 × 1024 ≈ 1.07e9 FLOPs per call. If the GPU run finishes in half the time, the FLOP rate roughly doubles; if it does not, the tensors are too small or the host is doing the work.

In code

import time, tensorflow as tf

print(tf.config.list_physical_devices("GPU"))

a = tf.random.normal((512, 1024))
b = tf.random.normal((1024, 1024))

def bench(dev, iters=200):
    with tf.device(dev):
        c = tf.matmul(a, b)                     # warm up
        t0 = time.perf_counter()
        for _ in range(iters):
            c = tf.matmul(a, b)
        c.numpy()                               # synchronize
        return (time.perf_counter() - t0) / iters

print("cpu", bench("/CPU:0"))
print("gpu", bench("/GPU:0"))

Check yourself

1. Why is .numpy() required before stopping the timer?
2. Why does a small batch size shrink the GPU advantage?
3. What does a CPU-bound tf.data pipeline look like in a profiler, and why does it hide GPU gains?

Key takeaways

• GPUs win on large, parallel tensor ops, not on small sequential work.
• Measure end to end with synchronization; kernel launch time is not execution time.
• Feeding the device fast enough is often the real bottleneck, not the compute.