Introduction

Using tf.data.Dataset for prediction is a good way to keep inference input handling aligned with the same pipeline style used for training. The main rule is simple: for model.predict, the dataset should usually yield features only, not (features, labels) pairs.

Build a Prediction Dataset

A common starting point is to build a dataset from NumPy arrays or tensors, then batch and prefetch it.

1import numpy as np
2import tensorflow as tf
3
4x_pred = np.random.rand(12, 4).astype("float32")
5
6pred_ds = (
7    tf.data.Dataset.from_tensor_slices(x_pred)
8    .batch(4)
9    .prefetch(tf.data.AUTOTUNE)
10)

Batching matters because Keras models expect batch-shaped input during prediction just as they do during training.

Train a Small Model and Predict

1import numpy as np
2import tensorflow as tf
3
4x_train = np.random.rand(100, 4).astype("float32")
5y_train = (x_train.sum(axis=1, keepdims=True) > 2).astype("float32")
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Input(shape=(4,)),
9    tf.keras.layers.Dense(16, activation="relu"),
10    tf.keras.layers.Dense(1, activation="sigmoid"),
11])
12
13model.compile(optimizer="adam", loss="binary_crossentropy")
14model.fit(x_train, y_train, epochs=3, batch_size=16, verbose=0)
15
16predictions = model.predict(pred_ds, verbose=0)
17print(predictions.shape)

That is the basic pattern: create the dataset, batch it, and pass it directly to predict.

Reuse a Labeled Dataset by Mapping Away the Labels

Sometimes the pipeline you already have yields (x, y) pairs because it was built for training or evaluation. Prediction should normally consume only the features.

1train_like_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(8)
2features_only_ds = train_like_ds.map(lambda x, y: x)
3
4predictions = model.predict(features_only_ds, verbose=0)
5print(predictions[:2])

This is often the cleanest way to reuse an existing pipeline without rebuilding it from scratch.

Keep IDs if Predictions Must Be Joined Back

Sometimes the output needs to be matched back to source records. In that case, include an identifier in the dataset and run inference with the model directly in a small loop.

1ids = np.arange(len(x_pred))
2id_ds = tf.data.Dataset.from_tensor_slices((ids, x_pred)).batch(4)
3
4rows = []
5for batch_ids, batch_x in id_ds:
6    batch_pred = model(batch_x, training=False).numpy().reshape(-1)
7    for row_id, score in zip(batch_ids.numpy(), batch_pred):
8        rows.append((int(row_id), float(score)))
9
10print(rows[:3])

This gives you tighter control than predict when ordering, IDs, or custom post-processing matter.

Why batch and prefetch Still Matter at Inference Time

People often think pipeline tuning matters only during training. In practice, inference pipelines benefit too.

• 'batch improves throughput and keeps shapes predictable.'
• 'prefetch overlaps input preparation with model execution.'
• stable ordering makes it easier to reconcile predictions with source rows.

For large prediction jobs, the input pipeline can become a meaningful part of the total runtime.

When to Call the Model Directly Instead of predict

model.predict is convenient when you just want batched outputs. Calling model(batch_x, training=False) directly is better when you need:

• custom output packaging,
• record IDs in the result,
• incremental writes,
• tighter control over the loop.

Both approaches are valid. The choice depends on how much control the inference pipeline needs.

Common Pitfalls

A common mistake is feeding a dataset that yields (x, y) pairs directly to predict and then being surprised by shape or structure issues. Prediction usually wants only x.

Another issue is forgetting to batch the dataset. That can hurt performance and also produce shapes that the model did not expect.

Teams also sometimes shuffle inference data and then struggle to match predictions back to the original records.

Summary

• 'tf.data.Dataset works well as an input pipeline for Keras prediction.'
• For model.predict, the dataset should usually yield features only.
• Batch and prefetch the dataset for reliable shapes and better throughput.
• Map (x, y) datasets to x when reusing training-style pipelines for inference.
• Include IDs or call the model directly when predictions must be joined back to source rows.