TensorFlow & Keras: Production Engine

Keras's .fit() is a very quick way to start: three lines and you're training. But most production projects soon need more than it offers: Multi-input models with mixed loss weights. Streaming data from S3 with on-the-fly augmentation. Mixed-precision training to fit a model on a single GPU. Custom training loops that alternate generator and discriminator updates. None of that is a simple model.fit(X, y).

This lesson takes you from "I trained a small Keras example" to "I released a TF/Keras model that handles real traffic." We'll cover the TF 2.x execution model, tf.data pipelines that keep the GPU busy, the three Keras APIs and when to use each, custom training with tf.GradientTape, useful callbacks, mixed precision, distribution strategies, and the saving formats you will use for deployment.

Run in Colab or locally with pip install tensorflow. Expected outputs are shown in comments. TensorFlow has no WebAssembly build, so unlike most of this site it cannot run in your browser. If you want to see a network train here, build one from scratch in NumPy instead. It reaches 97% on hand-written digits in about a second.

1. TF 2.x: Eager by Default, Graphs When You Ask

TF 1.x was graph-first: you defined a static computation graph, then ran it in a Session. Debugging was hard, because you could only inspect placeholders, not values. TF 2.x changed that: eager execution is the default, as in NumPy.

python

import tensorflow as tf

a = tf.constant([1.0, 2.0, 3.0])
b = tf.constant([4.0, 5.0, 6.0])
print(a + b)                # tf.Tensor([5. 7. 9.], shape=(3,), dtype=float32)
print((a + b).numpy())      # [5. 7. 9.]   (back to a NumPy array)

You can print(), set breakpoints, use pdb. Each op runs immediately. The cost: Python is slow, and running each op separately on the GPU is less efficient than running them combined.

tf.function solves this. Decorate a Python function, and TF traces it into a graph the first time it's called, then runs the compiled graph on every later call.

python

@tf.function
def fast_step(x, y):
    return tf.reduce_mean(tf.square(x - y))

# First call: traces the graph (slow).
# Later calls: the graph runs (fast, with fused operations and no Python overhead).
fast_step(a, b)

+ setup added so this can run · defines tf, a, b

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

tf = _AutoMock('tf')
a = _AutoMock('a')
b = _AutoMock('b')

Inside a tf.function-decorated function, you write Python, and TF traces it once and compiles it. Conditionals and loops that depend on tensor values are converted to tf.cond and tf.while_loop. Python print is not converted, so use tf.print for output during training (see Common Mistake 3). Code that runs on every training step should be inside @tf.function. Setup code that runs once can stay eager.

Note: tracing happens once per input signature. If you call fast_step with different shapes or dtypes, TF re-traces. Use tf.function(input_signature=[...]) to lock the signature in production.

2. Tensors vs NumPy: The Boundary

tf.Tensor and np.ndarray look similar, but a tensor can live on the GPU, while a NumPy array is always in CPU memory.

python

import numpy as np
import tensorflow as tf

np_arr = np.array([[1, 2], [3, 4]], dtype=np.float32)
tf_t   = tf.convert_to_tensor(np_arr)               # copies the data from the CPU to the device

print(tf_t.device)          # /job:localhost/replica:0/task:0/device:GPU:0  (if available)
print(tf_t.numpy())         # back to NumPy, blocking copy

Conversion costs:

NumPy → tf.Tensor copies data to the device (GPU/TPU). This takes time.
tf.Tensor → NumPy with .numpy() blocks until the op completes, then copies the data back to the CPU. While it waits, the CPU and GPU can't work in parallel.

Rule: keep everything as tensors for the whole training step. Only call .numpy() for logging, plotting, or returning final results. Calling .numpy() inside a training loop is a common reason for low GPU utilisation (for example, 30%).

3. `tf.data`: The Pipeline That Feeds the GPU

GPU time is expensive. If the GPU is waiting for data, you are paying for nothing. tf.data is the official pipeline API: it builds an asynchronous, parallel, prefetched stream of batches that overlaps preprocessing with model execution.

python

import tensorflow as tf

# 1. Source: where does data come from?
ds = tf.data.Dataset.from_tensor_slices((X_train, y_train))     # in-memory

# Or from files:
# ds = tf.data.Dataset.list_files("s3://bucket/images/*.jpg")

# 2. Shuffle (buffer large enough for true randomness)
ds = ds.shuffle(buffer_size=10_000, seed=42, reshuffle_each_iteration=True)

# 3. Map: preprocess in parallel across CPU threads
def preprocess(x, y):
    x = tf.cast(x, tf.float32) / 255.0
    return x, y

ds = ds.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)

# 4. Batch
ds = ds.batch(64, drop_remainder=True)

# 5. Prefetch: GPU works on batch N while CPU prepares batch N+1
ds = ds.prefetch(tf.data.AUTOTUNE)

# 6. Hand to model.fit
model.fit(ds, epochs=10)

+ setup added so this can run · defines X_train, y_train, model

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

X_train = _AutoMock('X_train')
y_train = _AutoMock('y_train')
model = _AutoMock('model')

Each step matters:

shuffle(buffer_size) must be larger than your batch size, ideally close to dataset size. If the buffer is too small, examples that are close together in the file end up in the same batches, so the data is not well mixed.
map(... , num_parallel_calls=AUTOTUNE) preprocesses in parallel. Without num_parallel_calls, it runs one element at a time, and single-threaded preprocessing is often the bottleneck.
batch before prefetch: prefetching whole batches also hides the time spent building each batch.
prefetch(AUTOTUNE) at the end lets TF decide how many batches to buffer. It is often the most important line for GPU utilisation.

Order matters. shuffle().batch() shuffles examples and then groups them, which is what you usually want. batch().shuffle() shuffles whole batches but keeps the contents of each batch fixed, which is almost never what you want.

For image pipelines, decoding and resizing belong inside .map():

python

def load_image(path):
    raw = tf.io.read_file(path)
    img = tf.io.decode_jpeg(raw, channels=3)
    img = tf.image.resize(img, [224, 224])
    return tf.cast(img, tf.float32) / 255.0

ds = tf.data.Dataset.list_files("data/*.jpg")
ds = ds.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.batch(32).prefetch(tf.data.AUTOTUNE)

+ setup added so this can run · defines tf

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

tf = _AutoMock('tf')

tf.io.decode_jpeg runs on CPU, in parallel, while the GPU trains on the previous batch. With a well-tuned pipeline, GPU utilisation can often reach 90% or more.

4. The Functional API: Beyond Sequential

Sequential is fine for linear stacks. As soon as you have two inputs, a branch, a shared layer, or multiple outputs, you need the Functional API.

python

from tensorflow import keras
from tensorflow.keras import layers

# Two-input multi-task model: image and table features in, a class and a price out
image_in = keras.Input(shape=(224, 224, 3), name="image")
tab_in   = keras.Input(shape=(10,),         name="tabular")

# Image branch
x = layers.Conv2D(32, 3, activation="relu")(image_in)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(64, activation="relu")(x)

# Tabular branch
t = layers.Dense(32, activation="relu")(tab_in)

# Merge
merged = layers.concatenate([x, t])
merged = layers.Dense(64, activation="relu")(merged)

# Two heads (multi-task)
class_out = layers.Dense(5, activation="softmax", name="class")(merged)
price_out = layers.Dense(1, name="price")(merged)

model = keras.Model(
    inputs=[image_in, tab_in],
    outputs={"class": class_out, "price": price_out},
)

model.compile(
    optimizer="adamw",
    loss={"class": "sparse_categorical_crossentropy", "price": "mse"},
    loss_weights={"class": 1.0, "price": 0.1},      # weight one loss against the other
    metrics={"class": ["accuracy"], "price": ["mae"]},
)

Named inputs and outputs let you pass dicts to .fit():

python

model.fit(
    {"image": img_array, "tabular": tab_array},
    {"class": class_labels, "price": prices},
    epochs=10,
)

+ setup added so this can run · defines model, img_array, tab_array, class_labels, prices

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

model = _AutoMock('model')
img_array = _AutoMock('img_array')
tab_array = _AutoMock('tab_array')
class_labels = _AutoMock('class_labels')
prices = _AutoMock('prices')

This is the API you will use most often in production. Shared layers work by reusing a layer object:

python

shared_embed = layers.Embedding(vocab_size, 64)
left  = shared_embed(left_in)
right = shared_embed(right_in)        # SAME weights (this is called a siamese network)

+ setup added so this can run · defines vocab_size, left_in, right_in, layers

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

vocab_size = _AutoMock('vocab_size')
left_in = _AutoMock('left_in')
right_in = _AutoMock('right_in')
layers = _AutoMock('layers')

5. Subclassing: Maximum Flexibility

When the model needs dynamic control flow, custom training behaviour, or unusual research ideas, subclass keras.Model:

python

class ResidualMLP(keras.Model):
    def __init__(self, hidden=128, num_classes=10):
        super().__init__()
        self.dense1 = layers.Dense(hidden, activation="relu")
        self.dense2 = layers.Dense(hidden, activation="relu")
        self.proj   = layers.Dense(hidden)
        self.head   = layers.Dense(num_classes, activation="softmax")

    def call(self, x, training=False):
        residual = self.proj(x)
        h = self.dense1(x)
        if training:
            h = tf.nn.dropout(h, rate=0.3)        # Python if-logic that can change per call
        h = self.dense2(h)
        return self.head(h + residual)

model = ResidualMLP()
model.build(input_shape=(None, 50))               # define input shape
model.summary()

+ setup added so this can run · defines keras, layers, tf

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

keras = _AutoMock('keras')
layers = _AutoMock('layers')
tf = _AutoMock('tf')

The training flag is passed through automatically: Keras passes training=True during .fit() and training=False during .evaluate()/.predict(). Layers like Dropout and BatchNorm rely on it.

Subclassing is the right choice for:

Research code where the architecture evolves between runs
Models with input-dependent control flow (e.g. recursive networks)
Anything you want to inherit from and override

Trade-off: Subclassed models can't be serialised to the Keras JSON config format without get_config() boilerplate. Use the Functional API by default, and subclass when you need it.

6. Custom Layers

When you need a layer that Keras doesn't have, subclass keras.layers.Layer. The pattern is __init__ (config), build (lazy weight creation when input shape is known), call (forward).

python

class ScaledDotAttention(layers.Layer):
    def __init__(self, dim, **kwargs):
        super().__init__(**kwargs)
        self.dim = dim
        self.scale = dim ** -0.5

    def build(self, input_shape):
        self.wq = self.add_weight(name="wq", shape=(input_shape[-1], self.dim))
        self.wk = self.add_weight(name="wk", shape=(input_shape[-1], self.dim))
        self.wv = self.add_weight(name="wv", shape=(input_shape[-1], self.dim))

    def call(self, x):
        q = x @ self.wq
        k = x @ self.wk
        v = x @ self.wv
        attn = tf.nn.softmax(q @ tf.transpose(k, [0, 2, 1]) * self.scale, axis=-1)
        return attn @ v

    def get_config(self):
        return {**super().get_config(), "dim": self.dim}

+ setup added so this can run · defines layers, kwargs, tf

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

layers = _AutoMock('layers')
kwargs = _AutoMock('kwargs')
tf = _AutoMock('tf')

get_config() is what lets .keras serialisation rebuild your custom layer with the right arguments. Loading also needs Keras to know the class: register it with @keras.saving.register_keras_serializable(), or pass custom_objects={"ScaledDotAttention": ScaledDotAttention}.

7. Custom Training Loops with `tf.GradientTape`

.fit() works well until you need something it can't do directly: GAN training (alternating generator/discriminator), reinforcement learning, custom losses with extra heads, gradient accumulation, gradient penalties. For these, write a manual loop with tf.GradientTape:

python

optimizer = keras.optimizers.AdamW(learning_rate=1e-3)
loss_fn   = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

@tf.function                                            # compile the step
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss   = loss_fn(y, logits)
        # Add weight decay or regularisation here if you want
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

for epoch in range(EPOCHS):
    for x, y in train_ds:
        loss = train_step(x, y)
    print(f"epoch {epoch} loss={loss.numpy():.4f}")

+ setup added so this can run · defines tf, EPOCHS, train_ds, keras, model

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

tf = _AutoMock('tf')
EPOCHS = _AutoMock('EPOCHS')
train_ds = [("alpha", 1), ("beta", 2), ("gamma", 3)]
keras = _AutoMock('keras')
model = _AutoMock('model')

The tape records every op on a watched variable, and tape.gradient(loss, vars) returns the gradient of loss with respect to each variable. The @tf.function decorator compiles train_step. Without it, each step pays Python's overhead.

GAN training is the standard example where the default .fit() is not enough:

python

@tf.function
def gan_step(real_x):
    noise = tf.random.normal([BATCH, LATENT])
    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        fake_x   = generator(noise, training=True)
        real_out = discriminator(real_x, training=True)
        fake_out = discriminator(fake_x, training=True)
        g_loss = generator_loss(fake_out)
        d_loss = discriminator_loss(real_out, fake_out)

    g_grads = gen_tape.gradient(g_loss, generator.trainable_variables)
    d_grads = disc_tape.gradient(d_loss, discriminator.trainable_variables)
    g_opt.apply_gradients(zip(g_grads, generator.trainable_variables))
    d_opt.apply_gradients(zip(d_grads, discriminator.trainable_variables))
    return g_loss, d_loss

+ setup added so this can run · defines tf, BATCH, LATENT, generator, discriminator, generator_loss, discriminator_loss, g_opt, d_opt

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

tf = _AutoMock('tf')
BATCH = _AutoMock('BATCH')
LATENT = _AutoMock('LATENT')
generator = _AutoMock('generator')
discriminator = _AutoMock('discriminator')
def generator_loss(*_a, **_kw):
    print('-> generator_loss() called')
    return _AutoMock('generator_loss()')
def discriminator_loss(*_a, **_kw):
    print('-> discriminator_loss() called')
    return _AutoMock('discriminator_loss()')
g_opt = _AutoMock('g_opt')
d_opt = _AutoMock('d_opt')

Two tapes, two optimisers, one step. The default .fit() can't do this. (You can also put this code in an overridden train_step method and still call .fit().) Custom loops give you full control, but you have to write the loop yourself.

8. Callbacks: Hooks Into `.fit()`

Callbacks let you customise .fit() without writing a manual loop. The ones you will use most:

python

from tensorflow.keras import callbacks

cbs = [
    callbacks.ModelCheckpoint(
        "models/best.keras",
        monitor="val_loss",
        save_best_only=True,            # only overwrite if val_loss improves
        save_weights_only=False,        # save full model
    ),
    callbacks.EarlyStopping(
        monitor="val_loss",
        patience=5,                     # stop if no improvement for 5 epochs
        restore_best_weights=True,      # go back to the best weights, not the last ones
    ),
    callbacks.ReduceLROnPlateau(
        monitor="val_loss",
        factor=0.5,
        patience=2,                     # halve the learning rate after 2 epochs with no improvement
        min_lr=1e-6,
    ),
    callbacks.TensorBoard(log_dir="logs/run01"),
    callbacks.CSVLogger("logs/train.csv"),
]

model.fit(train_ds, validation_data=val_ds, epochs=100, callbacks=cbs)

+ setup added so this can run · defines train_ds, model, val_ds

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

train_ds = _AutoMock('train_ds')
model = _AutoMock('model')
val_ds = _AutoMock('val_ds')

Always use restore_best_weights=True on EarlyStopping in production. Without it, you keep the weights from the epoch where training stopped, which are often worse than the weights from the best epoch.

To write a custom callback, subclass Callback and define the methods you need. Available hooks include on_train_begin, on_epoch_begin/end, on_batch_begin/end, etc.

python

class GradientNormLogger(callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        # Caveat: this requires a custom training loop that exposes grad norms,
        # or you can probe model.optimizer.iterations and so on, in on_epoch_end.
        if batch % 100 == 0:
            print(f"batch {batch} loss={logs['loss']:.4f}")

+ setup added so this can run · defines callbacks

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

callbacks = _AutoMock('callbacks')

9. Mixed Precision Training

NVIDIA GPUs from the Volta generation onwards have tensor cores that run float16 matmuls much faster than float32 (bfloat16 needs Ampere or later), and these values use half the memory. Mixed precision keeps the weights in float32 (for stable updates) but does the computation in float16. You turn it on with one line:

python

from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy("mixed_float16")    # weights stay float32, maths runs in float16

# Build the model AFTER setting the policy.
model = build_model()

+ setup added so this can run · defines build_model

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

def build_model(*_a, **_kw):
    print('-> build_model() called')
    return _AutoMock('build_model()')

One subtle requirement: the output layer should be float32 to avoid numerical issues in the loss:

python

outputs = layers.Dense(num_classes, dtype="float32")(x)    # force fp32 output

+ setup added so this can run · defines x, num_classes, layers

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

x = _AutoMock('x')
num_classes = _AutoMock('num_classes')
layers = _AutoMock('layers')

On the right hardware, mixed precision often gives about a 2× speedup and lets you fit a larger batch in the same memory, with no other code changes. On older GPUs (Pascal and earlier) without tensor cores, it gives little or no speedup.

bfloat16 (Brain Float 16) has the dynamic range of float32 but less precision than float16 (8 significant bits vs 11). The wide range means far fewer overflow problems, so it needs no loss scaling. Available on TPUs and Ampere+ GPUs. Use "mixed_bfloat16" on TPUs.

10. Distributed Training

If one GPU isn't enough, tf.distribute.MirroredStrategy replicates the model across all local GPUs and averages gradients:

python

strategy = tf.distribute.MirroredStrategy()
print(f"Replicas: {strategy.num_replicas_in_sync}")    # for example 4

with strategy.scope():
    model = build_model()                              # weights created on each GPU
    model.compile(optimizer="adamw", loss="...")

model.fit(train_ds, epochs=10)                         # data sharded automatically

+ setup added so this can run · defines train_ds, build_model, tf

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

train_ds = _AutoMock('train_ds')
def build_model(*_a, **_kw):
    print('-> build_model() called')
    return _AutoMock('build_model()')
tf = _AutoMock('tf')

The with strategy.scope() block is required: the model must be created inside it, so the weights are mirrored. Multi-host: MultiWorkerMirroredStrategy. TPUs: TPUStrategy. The Keras API is identical. Only the strategy object changes.

Effective batch size scales with replica count. If you used batch=64 on one GPU and now run on 4 GPUs, the effective batch is 256. You usually scale the LR by the same factor (lr * 4) to compensate, and add a warmup.

TPUs (Google Cloud) are very good for training large transformers, but they have some problems: you need bfloat16 for good performance, dynamic shapes are slow, and debugging is harder. Use them only when your workload is big enough to need them.

11. Saving Formats: How You Will Deploy

There are four formats you will meet when you save and deploy a Keras model. Pick the right one for the deployment target.

Format	File / Dir	Use case
`.keras`	single ZIP file	Default. Architecture + weights + optimiser state. Use for checkpointing and most reloading.
SavedModel	directory tree	TensorFlow Serving, TF.js conversion, TFLite conversion. Contains the graph, weights, and signature defs.
TFLite	`.tflite` file	Mobile and edge (Android, iOS, embedded). Quantised, optimised for small runtime.
TF.js	JSON + binary shards	Browser inference. Convert from SavedModel via `tensorflowjs_converter`.

python

# .keras format: best when you reload the model in Python
model.save("models/sentiment.keras")
loaded = keras.models.load_model("models/sentiment.keras")

# SavedModel format: for TensorFlow Serving
model.export("models/sentiment_savedmodel")          # Keras 3 API

# TFLite format: for mobile apps
converter = tf.lite.TFLiteConverter.from_saved_model("models/sentiment_savedmodel")
tflite_bytes = converter.convert()
with open("sentiment.tflite", "wb") as f:
    f.write(tflite_bytes)

+ setup added so this can run · defines model, keras, tf

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

model = _AutoMock('model')
keras = _AutoMock('keras')
tf = _AutoMock('tf')

save_weights_only=True in ModelCheckpoint saves only the weights, not the architecture. That is fine if you can rebuild exactly the same model (from the same code), but risky otherwise. For a model you will reload after the code has changed, save the full model.

12. Keras vs PyTorch: When to Use Each

Dimension	Keras (TF)	PyTorch
Default user	Production engineers, applied ML, fast prototypes	Researchers, paper authors
API style	Declarative (`.compile`, `.fit`) plus subclassing	Imperative, more "just Python"
Production deployment	TF Serving, TFLite, TF.js: mature tools	TorchServe, ONNX, mobile via Torch Mobile
Distributed training	`tf.distribute` strategies: a clean abstraction	`torch.distributed` / DDP: more setup code
Research ecosystem	Smaller: most new papers release PyTorch code first	Dominant: HuggingFace, timm, MMDetection
Static graph optimisation	`tf.function` (opt-in)	`torch.compile` (added in PyTorch 2.0)
TPU support	First-class	Newer, less mature

In 2026 the differences are smaller than they were in 2020, because each framework has added the other's strengths. A practical rule: if you're building a product to deploy, use Keras/TF. If you're doing research and reimplementing papers, use PyTorch. Many teams use one framework in production and let researchers prototype in either, with a conversion step (often through ONNX) between the two.

Common Mistakes

1. Running each training step eagerly in a Python loop

python

for x, y in train_ds:        # fine in a Python loop?
    grads = compute(x, y)    # slow if compute() is not a tf.function

+ setup added so this can run · defines train_ds, compute

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

train_ds = [("alpha", 1), ("beta", 2), ("gamma", 3)]
def compute(*_a, **_kw):
    print('-> compute() called')
    return _AutoMock('compute()')

Iterating a tf.data.Dataset in a Python for loop is correct, and prefetch still prepares the next batches in the background. The slow part is the step itself: if compute runs eagerly, every step pays Python overhead. Wrap the training step in @tf.function (the official custom training loop guide does this), or pass the dataset directly to .fit().

2. Not using tf.function on hot paths

Your train_step runs 100,000 times. Without @tf.function, each call pays Python's overhead, and it is often several times slower than the compiled version. Profile before you assume. But on code that runs every step, the decorator is an easy speedup.

3. Mixing eager and graph code

Inside a @tf.function-traced function, Python print() runs only during tracing, not on later calls. Use tf.print() for runtime output. Same trap with if over tensor values (use tf.cond), Python for over tensors (use tf.while_loop or vectorise), and list appending (use tf.TensorArray).

4. Saving weights when you needed architecture too

python

model.save_weights("ckpt.weights.h5")    # weights only, so you must rebuild the model to load them
# 3 months later: which version of the code defined this architecture?

+ setup added so this can run · defines model

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

model = _AutoMock('model')

Save the full model (.keras) unless you have a specific reason for weights-only. The extra disk space is worth it, because you can reload the model without searching old code for the architecture.

5. Forgetting to set_global_policy before building the model

python

model = build_model()                         # weights created in fp32
mixed_precision.set_global_policy("mixed_float16")    # too late

+ setup added so this can run · defines build_model, mixed_precision

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

def build_model(*_a, **_kw):
    print('-> build_model() called')
    return _AutoMock('build_model()')
mixed_precision = _AutoMock('mixed_precision')

The policy is read when each layer is created. Set it first, then build. Distribution strategies have the same trap: with strategy.scope(): must wrap the model construction.

6. Calling .numpy() inside the training loop

Even for logging, calling .numpy() every step forces a host-device sync. Either log every N steps, or accumulate with a tf.keras.metrics.Mean() object inside the graph and read it back once per epoch.

🎯 Your Turn: A Full `tf.data` Pipeline

Build a tf.data pipeline that takes a list of image file paths and corresponding integer labels, and produces shuffled, batched, prefetched (image, label) pairs ready for .fit().

Requirements:

Accept file_paths (list of strings) and labels (list of ints).
Each step in the pipeline:

- Read JPEG bytes from disk. - Decode to a 3-channel uint8 tensor. - Resize to (224, 224). - Cast to float32 and scale to [0, 1].

Shuffle with a buffer of 1000.
Batch into 32, dropping the remainder.
Prefetch with AUTOTUNE.
Use num_parallel_calls=AUTOTUNE on the decoding map.

python

import tensorflow as tf

def build_pipeline(file_paths, labels, batch_size=32):
    # TODO 1: create a Dataset from the (paths, labels) pair
    # TODO 2: shuffle with buffer 1000
    # TODO 3: map a function that loads + decodes + resizes + normalises
    #         use num_parallel_calls=tf.data.AUTOTUNE
    # TODO 4: batch with drop_remainder=True
    # TODO 5: prefetch with AUTOTUNE
    ...

# Usage
paths  = ["data/cat1.jpg", "data/dog1.jpg", "data/cat2.jpg", ...]
labels = [0, 1, 0, ...]
ds = build_pipeline(paths, labels, batch_size=32)

for images, lbls in ds.take(1):
    print(images.shape, images.dtype)    # (32, 224, 224, 3) float32
    print(lbls.shape, lbls.dtype)        # (32,) int32

Hint 1: A function that takes a path and label, returns image and label

tf.data.Dataset.map calls your function on each element. If the dataset yields (path, label) pairs, the mapped function signature is def load(path, label):. Decode the image inside, and return (image, label) unchanged. The label passes through untouched.

Hint 2: Decoding JPEGs in graph mode

Use tf.io.read_file(path) to read bytes from disk (works inside tf.data), then tf.io.decode_jpeg(bytes, channels=3) to get a uint8 tensor of shape (H, W, 3). Then tf.image.resize for the resize (returns float32 automatically) and divide by 255.0 for the normalisation. Don't use PIL or OpenCV inside the map: TensorFlow can't trace their calls into the graph.

Show full solution

python

import tensorflow as tf

AUTOTUNE = tf.data.AUTOTUNE

def load_and_preprocess(path, label):
    """Read JPEG bytes, decode, resize, normalise."""
    raw   = tf.io.read_file(path)
    image = tf.io.decode_jpeg(raw, channels=3)          # uint8 (H, W, 3)
    image = tf.image.resize(image, [224, 224])          # returns float32 (the next line makes the cast explicit)
    image = tf.cast(image, tf.float32) / 255.0          # explicit cast, then scale to [0, 1]
    return image, label

def build_pipeline(file_paths, labels, batch_size=32):
    ds = tf.data.Dataset.from_tensor_slices((file_paths, labels))
    ds = ds.shuffle(buffer_size=1000, seed=42, reshuffle_each_iteration=True)
    ds = ds.map(load_and_preprocess, num_parallel_calls=AUTOTUNE)
    ds = ds.batch(batch_size, drop_remainder=True)
    ds = ds.prefetch(AUTOTUNE)
    return ds


# Demo with synthetic paths/labels (requires real JPEGs on disk to actually iterate)
paths  = tf.constant(["/path/to/img1.jpg", "/path/to/img2.jpg"])
labels = tf.constant([0, 1], dtype=tf.int32)
ds = build_pipeline(paths, labels, batch_size=2)

# Expected output on real data:
# (2, 224, 224, 3) float32
# (2,) int32

What makes this pipeline fast:

num_parallel_calls=AUTOTUNE runs JPEG decoding on several CPU threads in parallel. Image decoding is your main CPU cost, and each file can be decoded independently.
prefetch(AUTOTUNE) lets TF buffer N batches ahead while the GPU is busy training on the current one. AUTOTUNE picks N adaptively based on observed throughput.
drop_remainder=True keeps the batch size fixed. This matters for @tf.function tracing (a changing batch size causes re-tracing), and TPUs require it.
shuffle before map before batch: shuffling file paths is cheap (they are short strings), but shuffling decoded images would waste memory and time, because each image is a large tensor.

Additions you will often need in production:

Caching: .cache() after .map() if the preprocessed data fits in memory. The next epoch then reads from RAM instead of disk.
Augmentation: a second .map() after batching with tf.image.random_flip_left_right, tf.image.random_crop, etc. This runs on the CPU in the pipeline. Or move it into a Keras preprocessing layer that runs on the GPU.
Sharded reading from cloud: tf.data.Dataset.list_files("gs://bucket/*.tfrecord") and tf.data.TFRecordDataset(...) for cloud-scale training.

If your GPU utilisation is below ~85% during training, look at the pipeline first. The cause is usually a missing num_parallel_calls, a missing prefetch, or a .map step that runs Python code without you noticing (PIL, OpenCV, slow custom ops).

What You Learned

TF 2.x is eager by default, and @tf.function traces a function into a graph for speed. Wrap hot paths.
tf.Tensor ↔ np.ndarray conversion costs a device sync. Keep tensors as tensors inside training loops, and call .numpy() only at the boundaries.
tf.data pipelines with shuffle → map(parallel) → batch → prefetch can raise GPU utilisation from around 30% to around 90%.
Functional API for multi-input, multi-output, shared, and branched models. Sequential only for linear stacks. Subclassing for dynamic control flow in research code.
Custom layers subclass keras.layers.Layer, override build and call. get_config() enables round-trip serialisation.
tf.GradientTape for custom training loops: GANs, RL, gradient accumulation, and anything the default .fit() can't do.
Callbacks: ModelCheckpoint(save_best_only=True), EarlyStopping(restore_best_weights=True), ReduceLROnPlateau, TensorBoard. Always restore best weights.
Mixed precision (mixed_float16) often gives about a 2× speedup on GPUs with tensor cores. Force the output layer to float32.
MirroredStrategy for multi-GPU: wrap model construction in strategy.scope(). Scale LR with replica count.
Save as .keras for Python reload, SavedModel for serving, TFLite for mobile, TF.js for browser.
Keras vs PyTorch: Keras for production and fast prototypes, PyTorch for research. Both are good choices in 2026.

Next: Convolutional Neural Networks: the architecture family behind many modern vision systems, with the maths, the architectures, and modern techniques.

1. TF 2.x: Eager by Default, Graphs When You Ask

2. Tensors vs NumPy: The Boundary

3. tf.data: The Pipeline That Feeds the GPU

4. The Functional API: Beyond Sequential

5. Subclassing: Maximum Flexibility

6. Custom Layers

7. Custom Training Loops with tf.GradientTape

8. Callbacks: Hooks Into .fit()

9. Mixed Precision Training

10. Distributed Training

11. Saving Formats: How You Will Deploy

12. Keras vs PyTorch: When to Use Each

Common Mistakes

🎯 Your Turn: A Full tf.data Pipeline

What You Learned

3. `tf.data`: The Pipeline That Feeds the GPU

7. Custom Training Loops with `tf.GradientTape`

8. Callbacks: Hooks Into `.fit()`

🎯 Your Turn: A Full `tf.data` Pipeline