The Eager / Graph / @tf.function Mental Model

The single most important model to lock in your head

TF 2.x runs in two modes you switch between depending on goal: eager for debugging, graph for speed. The decorator that flips between them is @tf.function. Get this mental model right now and the entire rest of TF makes sense.

Eager mode: every op runs immediately, returns a concrete tensor, and you can call .numpy() on it. print(x) works. pdb works. x.shape reflects the actual data. This is the default in TF 2.x and the only mode that really feels like Python.

Graph mode: when you wrap a function in @tf.function, TF traces the first call — running it with symbolic placeholder tensors and recording a static computation graph. From the second call onward, that compiled graph runs directly on the runtime, skipping Python overhead and enabling op fusion plus XLA. Speedups of 2–10× are typical for tight inner loops.

The trap: print() inside a @tf.function-decorated function only fires during tracing, not on every call. If your prints suddenly stopped, that's why. Use tf.print() instead — it's part of the traced graph.

Code

Eager — NumPy랑 똑같은 느낌·python

import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0])
print(x + 1)  # tf.Tensor([2. 3. 4.], shape=(3,), dtype=float32)
print(tf.executing_eagerly())  # True

# Mix freely with NumPy
import numpy as np
print(np.mean(x))   # 2.0
print(x.numpy())    # array([1., 2., 3.], dtype=float32)

@tf.function — 같은 코드, graph로 컴파일·python

import tensorflow as tf

@tf.function
def train_step(x, y, model, optimizer, loss_fn):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
    grads = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    return loss

# Decorate ONCE outside the loop. Inside the loop = retrace per call.
for x_batch, y_batch in train_ds:
    loss = train_step(x_batch, y_batch, model, optimizer, loss_fn)

print()는 거짓말, tf.print()는 진실·python

@tf.function
def step(x):
    print('python print:', x)  # only on tracing, NOT every call
    tf.print('tf print:', x)   # fires every call, prints to stderr
    return x + 1

for i in range(3):
    step(tf.constant(float(i)))

The Eager / Graph / @tf.function Mental Model

The single most important model to lock in your head

Code

External links

Exercise

Progress

Comments 0