Deep Learning from Scratch, Part 20: Tensors, Mini-Batches, and tf.data Pipelines

In Part 19, we trained a Keras neural network by passing NumPy arrays directly to model.fit().

That approach is perfectly valid for a small dataset that fits comfortably in memory. But as data preparation becomes more complex, the model needs more than a collection of arrays. It needs an organized route through which observations are read, transformed, shuffled, grouped, and delivered during training.

A data pipeline is simply that route: the sequence of operations that carries raw observations to the model in a form it can process efficiently.

In this article, we will build a TensorFlow tf.data pipeline for structured data. Along the way, we will examine what tensors actually represent, why mini-batches affect optimization, how shuffling works, and what caching and prefetching can—and cannot—improve.

The important distinction is this:

A data pipeline can improve how efficiently data reaches the model. It does not automatically improve the model’s statistical accuracy.

From NumPy Arrays to Tensors

In NumPy, our structured input data was represented by a two-dimensional array:

• Each row represented one observation.
• Each column represented one feature.
• A separate one-dimensional array stored the labels.

TensorFlow expresses the same numerical structure using tensors.

A tensor is a multidimensional array whose elements share a common data type. Like a NumPy array, it has a shape and a data type. Unlike a typical NumPy array, a TensorFlow tensor is immutable: operations create new tensors instead of changing the existing tensor’s contents. TensorFlow tensor guide

The word tensor can sound more mysterious than the object itself. A scalar is a rank-zero tensor, a vector is a rank-one tensor, and a matrix is a rank-two tensor. Higher-rank tensors simply add more axes.

For structured data, our features usually form a rank-two tensor.

import tensorflow as tf
features_tensor = tf.convert_to_tensor(
    X_train_scaled,
    dtype=tf.float32,
)
labels_tensor = tf.convert_to_tensor(
    y_train,
    dtype=tf.float32,
)
print(features_tensor.shape)
print(features_tensor.dtype)
print(labels_tensor.shape)
print(labels_tensor.dtype)

If the training set contains 460 observations and eight features, the feature tensor has shape (460, 8). The label tensor has shape (460,).

The number 460 is not built into the network architecture. It is simply the number of observations currently stored in this tensor. The model’s input definition still describes one observation at a time:

keras.Input(shape=(8,))

The missing batch dimension is added when observations are grouped for training.

[Insert Figure 1 here]

Shape, Rank, and Data Type Answer Different Questions

These three tensor properties are related but not interchangeable.

Shape

Shape tells us how many elements exist along each axis.

For a structured-data batch with 32 observations and eight features, the input shape is (32, 8).

Rank

Rank is the number of axes.

A tensor with shape (32, 8) has rank two. A tensor with shape (32,) has rank one.

Rank is not the number of observations, nor is it the number of features.

Data Type

The data type tells TensorFlow how each numerical value is stored and interpreted.

Neural networks commonly use float32 because it offers a practical balance between precision, memory use, and computational speed. Accidentally supplying float64, integers, or strings may cause conversions, increased memory use, or errors in later operations.

print("Feature rank:", tf.rank(features_tensor).numpy())
print("Feature shape:", features_tensor.shape)
print("Feature dtype:", features_tensor.dtype)

Shape describes structure. Rank describes dimensionality. Data type describes representation.

A valid pipeline must preserve all three in the form expected by the model.

Why We Train with Mini-Batches

In full-batch gradient descent, the model processes every training observation before performing one parameter update.

In stochastic gradient descent at its most literal extreme, the model updates its parameters after one observation.

Mini-batch training operates between these two cases. It calculates an update from a subset of observations, such as 16, 32, or 64 examples.

If the training set contains 460 observations and the batch size is 32, the model receives 14 complete batches and one smaller final batch containing 12 observations.

Unless we explicitly request otherwise, TensorFlow keeps that final incomplete batch. This is usually the sensible choice for a small structured dataset because every training example remains available.

Mini-batches are not used only because the complete dataset may be too large for memory. They also change the optimization process.

A Mini-Batch Gradient Is an Estimate

The full training gradient averages information from every training observation. A mini-batch gradient uses only the observations in the current batch.

This means a mini-batch update is a noisy estimate of the full-batch update.

That noise is not necessarily a defect. It can help the optimizer move through complicated regions of the loss surface rather than following an identical, perfectly averaged direction at every step.

However, more noise is not always better.

Smaller Batches

Smaller batches generally provide:

• More parameter updates per epoch
• Noisier gradient estimates
• Lower memory requirements
• Less efficient hardware utilization when made extremely small

Larger Batches

Larger batches generally provide:

• Fewer parameter updates per epoch
• More stable gradient estimates
• Higher memory requirements
• Better vectorization up to the limits of the hardware

Batch size therefore affects both computation and optimization. It should be treated as a hyperparameter, not as a fixed property of structured data.

A batch size of 32 is a reasonable starting point for this example, but it is not universally optimal.

[Insert Figure 2 here]

What tf.data.Dataset Represents

A tf.data.Dataset represents a sequence of elements. An element may be a single tensor, a feature-label pair, a dictionary, or a more complex nested structure.

The pipeline is built by starting with a data source and applying transformations that return new datasets. TensorFlow’s official guide describes Dataset as the central abstraction for constructing reusable input pipelines. TensorFlow tf.data guide

For supervised learning, each element usually contains:

(features, label)

We can create this structure from our NumPy arrays:

raw_train_dataset = tf.data.Dataset.from_tensor_slices(
    (X_train_scaled, y_train)
)

from_tensor_slices() slices both arrays along their first axis. Observation zero is paired with label zero, observation one with label one, and so on.

This alignment is critical. Shuffling a feature array independently from its label array would destroy the relationship the model is expected to learn.

We can inspect the dataset’s element specification:

print(raw_train_dataset.element_spec)

For eight input features, the result should resemble:

(
    TensorSpec(shape=(8,), dtype=tf.float32),
    TensorSpec(shape=(), dtype=tf.float32)
)

At this stage, one dataset element contains one feature vector and one scalar label. The dataset has not yet been batched.

single_features, single_label = next(
    iter(raw_train_dataset)
)
print(single_features.shape)
print(single_label.shape)

The feature shape is (8,), while the scalar label has shape ().

Dataset Transformations Are a Recipe

Creating a Dataset does not normally execute the complete pipeline immediately. Transformations describe how elements should be produced when the dataset is iterated.

This is sometimes called lazy execution.

For example:

dataset = raw_train_dataset.shuffle(460)
dataset = dataset.batch(32)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

Each line produces a new dataset description. The observations begin flowing through those operations when training or explicit iteration requests them.

[Insert Figure 3 here]

The order of the transformations matters. These two pipelines are not equivalent:

dataset.shuffle(460).batch(32)

dataset.batch(32).shuffle(15)

The first shuffles individual observations and then forms batches.

The second forms batches first and then changes the order of those already-formed batches. The observations inside each batch remain grouped together.

For ordinary supervised training, we normally want to shuffle observations before batching.

Building a Reusable tf.data Function

We can collect the pipeline logic in one function:

import numpy as np
import tensorflow as tf
BATCH_SIZE = 32
SHUFFLE_SEED = 42
def make_dataset(
    features,
    labels,
    batch_size=BATCH_SIZE,
    training=False,
):
    features = np.asarray(
        features,
        dtype="float32",
    )
    labels = np.asarray(
        labels,
        dtype="float32",
    )
    dataset = tf.data.Dataset.from_tensor_slices(
        (features, labels)
    )
    if training:
        dataset = dataset.shuffle(
            buffer_size=len(features),
            seed=SHUFFLE_SEED,
            reshuffle_each_iteration=True,
        )
    dataset = dataset.batch(
        batch_size,
        drop_remainder=False,
    )
    dataset = dataset.prefetch(
        tf.data.AUTOTUNE
    )
    return dataset

We then create separate datasets for training, validation, and testing:

train_dataset = make_dataset(
    X_train_scaled,
    y_train,
    training=True,
)
validation_dataset = make_dataset(
    X_validation_scaled,
    y_validation,
    training=False,
)
test_dataset = make_dataset(
    X_test_scaled,
    y_test,
    training=False,
)

The three datasets do not have identical responsibilities.

• The training dataset is shuffled and used for parameter updates.
• The validation dataset is not used for parameter updates.
• The test dataset remains untouched until final evaluation.

The pipeline should preserve the experimental boundaries established in Part 18. tf.data organizes observations; it does not redefine which observations belong to training, validation, or testing.

Inspecting a Batch Before Training

Before giving a dataset to Keras, inspect one batch:

batch_features, batch_labels = next(
    iter(train_dataset)
)
print("Feature batch shape:", batch_features.shape)
print("Label batch shape:", batch_labels.shape)
print("Feature dtype:", batch_features.dtype)
print("Label dtype:", batch_labels.dtype)

For a batch size of 32 and eight features, the expected shapes are:

Feature batch shape: (32, 8)
Label batch shape: (32,)

This small inspection catches several common problems:

• The label column may still be present in the feature tensor.
• Features and labels may have different numbers of observations.
• The features may have an unexpected data type.
• A transformation may have removed or added an axis.
• The model may expect named inputs while the dataset supplies an unnamed tensor.

A pipeline that can be constructed is not necessarily a pipeline that contains what we intended.

How Shuffling Actually Works

Dataset.shuffle() maintains a buffer and samples the next element from that buffer. A larger buffer generally produces more thorough mixing but requires more memory. TensorFlow tf.data guide

For our small training set, we can use the complete training-set size:

buffer_size=len(X_train_scaled)

This permits an approximate full shuffle of the observations before batching.

For a dataset containing millions of observations, creating a full-size in-memory shuffle buffer may be impractical. A smaller buffer still provides local randomization, but it does not generate the same distribution of permutations as shuffling the complete dataset.

Why Shuffle the Training Set?

Suppose the source file is sorted by target class, time, hospital, geography, or acquisition device. Without shuffling, consecutive mini-batches may represent very different subpopulations.

The early updates might then be dominated by one region of the data, followed by updates dominated by another. The model would chase the ordering of the file instead of receiving reasonably mixed gradient estimates.

Why Not Shuffle Validation and Test Data?

Validation and test metrics generally do not require a random order because their parameters are not updated. A stable order also makes predictions easier to match back to their original rows during error analysis.

Shuffling validation data does not make the evaluation more honest. The important requirement is that the validation observations remain separate from training.

What Does reshuffle_each_iteration=True Do?

It requests a new shuffle order whenever the dataset begins a new iteration, such as a new training epoch.

Using the same random seed does not mean every epoch must have the same order. The seed makes the random process reproducible under compatible conditions, while reshuffle_each_iteration=True allows the order to change between iterations.

Exact reproducibility can still be affected by software versions, parallel execution, device operations, and nondeterministic kernels. A seed controls an important source of randomness, but it is not a universal guarantee that every hardware environment will produce identical final weights.

Batching and the Final Partial Batch

Our pipeline uses:

dataset.batch(
    BATCH_SIZE,
    drop_remainder=False,
)

With drop_remainder=False, the last batch may contain fewer observations than the requested batch size.

This is usually appropriate for validation and testing because discarding observations would change the evaluated dataset. It is also sensible for a small training set, where every example matters.

Setting drop_remainder=True gives every batch an identical first dimension, which can be useful for certain compiled or distributed workloads. But it also discards the final incomplete batch.

If the training dataset is reshuffled every epoch, the discarded observations may differ from epoch to epoch. That is not necessarily wrong, but it should be an intentional decision rather than an unnoticed side effect.

Prefetching: Preparing the Next Batch Early

Without prefetching, the training process may behave like this:

1. The model waits for the next batch.
2. The pipeline prepares the batch.
3. The model performs its computation.
4. The model waits again.

Prefetching allows data preparation and model computation to overlap. While the model processes the current batch, the pipeline can prepare a later batch in the background.

dataset = dataset.prefetch(
    tf.data.AUTOTUNE
)

tf.data.AUTOTUNE allows the runtime to choose the prefetch buffer dynamically. TensorFlow recommends prefetching when producer work can overlap with consumer work. TensorFlow data performance guide

[Insert Figure 4 here]

Prefetching can reduce the time the accelerator spends waiting for data. It does not change the network architecture, loss function, or mathematical definition of a gradient.

If the original pipeline is already extremely fast—as it may be for a tiny in-memory table—the measured speed difference may be negligible. The transformation becomes more valuable when reading files, performing expensive mappings, or feeding a faster GPU or TPU.

Caching: Useful, but Not Automatically Necessary

Caching stores dataset elements after certain operations have been executed. On later iterations, TensorFlow can reuse the cached results instead of repeating those earlier operations.

dataset = (
    source_dataset
    .map(
        expensive_preprocessing,
        num_parallel_calls=tf.data.AUTOTUNE,
    )
    .cache()
    .shuffle(shuffle_buffer)
    .batch(BATCH_SIZE)
    .prefetch(tf.data.AUTOTUNE)
)

Operations before cache() are reused from the cache. Operations after it are performed again during later iterations.

Placing shuffle() after cache() allows the cached observations to be reshuffled across epochs. Caching a shuffled dataset after the shuffle may unintentionally preserve one particular order, depending on how the pipeline is constructed and iterated.

TensorFlow’s performance guide recommends caching expensive deterministic transformations when the resulting data fits in memory or suitable local storage. TensorFlow data performance guide

However, our current arrays already fit in memory and require almost no per-example processing. Adding .cache() here may provide little practical benefit while making the example appear more complicated.

Caching is a tool for avoiding repeated work. It is not a mandatory decoration for every pipeline.

Mapping Transformations Across the Dataset

Dataset.map() applies a function to each dataset element.

For example, we might clip extreme standardized values:

def clip_features(features, label):
    clipped_features = tf.clip_by_value(
        features,
        clip_value_min=-5.0,
        clip_value_max=5.0,
    )
    return clipped_features, label
train_dataset = train_dataset.map(
    clip_features,
    num_parallel_calls=tf.data.AUTOTUNE,
)

This is only an implementation example, not a recommendation to clip every dataset at five standard deviations.

More importantly, pipeline transformations do not remove the risk of leakage.

A fixed transformation such as type casting or clipping does not learn population statistics. A stateful transformation such as standardization, vocabulary construction, or imputation does.

Any learned preprocessing state must be fitted or adapted using training data only.

Moving preprocessing into Dataset.map() does not make it safe if its statistics were calculated from the complete dataset.

Rebuilding the Same Keras Model

To isolate the effect of the input pipeline, we will keep the architecture from Part 19 unchanged.

import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
from keras import layers, regularizers
keras.utils.set_random_seed(42)
def build_model(number_of_features):
    model = keras.Sequential(
        [
            keras.Input(
                shape=(number_of_features,),
                name="features",
            ),
            layers.Dense(
                32,
                activation="relu",
                kernel_regularizer=regularizers.L2(1e-4),
                name="hidden_1",
            ),
            layers.Dropout(
                0.20,
                name="dropout",
            ),
            layers.Dense(
                16,
                activation="relu",
                kernel_regularizer=regularizers.L2(1e-4),
                name="hidden_2",
            ),
            layers.Dense(
                1,
                activation="sigmoid",
                name="probability",
            ),
        ],
        name="structured_data_classifier",
    )
    model.compile(
        optimizer=keras.optimizers.Adam(
            learning_rate=1e-3,
        ),
        loss=keras.losses.BinaryCrossentropy(),
        metrics=[
            keras.metrics.BinaryAccuracy(
                name="accuracy"
            ),
            keras.metrics.AUC(name="auc"),
        ],
    )
    return model
model = build_model(
    number_of_features=X_train_scaled.shape[1]
)

The architecture, optimizer, learning rate, loss, and regularization remain unchanged. Only the method of delivering data to the model is different.

This isolation matters. If we changed both the model and the pipeline, we could not attribute any observed difference to one specific cause.

Training Keras Directly from tf.data

Keras accepts a tf.data.Dataset directly in fit(), evaluate(), and predict(). When the dataset has already been batched, we should not pass another batch_size argument to fit(). Keras training and evaluation guide

early_stopping = keras.callbacks.EarlyStopping(
    monitor="val_loss",
    mode="min",
    patience=25,
    min_delta=1e-4,
    restore_best_weights=True,
)
history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=1000,
    callbacks=[early_stopping],
    verbose=0,
)

Compare this with Part 19:

model.fit(
    X_train_scaled,
    y_train,
    batch_size=32,
    ...
)

With NumPy input, Keras creates the batches.

With tf.data input, the dataset pipeline already defines them.

Do not specify batch_size=32 again and assume Keras will rebatch the dataset. Once a batched dataset is supplied, batching belongs to the pipeline.

Keras resets a finite dataset at the end of an epoch, so the same dataset object can be used across multiple epochs.

Why We Cannot Use validation_split Here

With NumPy arrays, Keras can create a validation subset by indexing the arrays.

A general Dataset may stream data from files, generators, services, or transformations without offering random sample indexing. Consequently, validation_split is not supported for Dataset inputs. Keras expects a separate validation dataset. Keras training and evaluation guide

This is not a limitation for our experiment because we already created explicit training and validation sets before scaling.

history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    ...
)

Explicit validation data also makes the experimental boundary easier to inspect.

[Insert Figure 5 here]

Counting Batches and Training Steps

We can inspect the number of batches:

number_of_training_batches = (
    tf.data.experimental.cardinality(
        train_dataset
    ).numpy()
)
print(
    "Training batches per epoch:",
    number_of_training_batches,
)

For a finite dataset, Keras can usually infer the number of steps per epoch from this cardinality.

A training step processes one batch and performs one optimizer update. Therefore, if an epoch contains 15 batches, the optimizer normally performs 15 updates during that epoch.

The number of epochs alone does not describe how many updates occurred.

Changing the batch size changes the number of updates per epoch. This is one reason two experiments with different batch sizes may not be directly comparable merely because they trained for the same number of epochs.

Be Careful with repeat()

The repeat() transformation can make a dataset produce multiple passes over its observations—or repeat indefinitely.

infinite_dataset = train_dataset.repeat()

An infinite dataset has no natural epoch boundary. If it is passed to fit(), we must define how many batches constitute one epoch:

model.fit(
    infinite_dataset,
    epochs=100,
    steps_per_epoch=number_of_training_batches,
)

Using repeat() without understanding steps_per_epoch can produce unexpected epoch definitions or training that does not terminate as intended.

For our finite structured dataset, repeat() is unnecessary. Keras already reinitializes the finite dataset for each epoch.

Does tf.data Improve Model Accuracy?

Not by itself.

If two training procedures use:

• The same observations
• The same feature values
• The same labels
• The same batch groupings and order
• The same initial parameters
• The same optimizer
• The same random operations

then replacing NumPy input with an equivalent tf.data pipeline does not introduce a new learning algorithm.

The main purpose of tf.data is to make input handling scalable, reusable, and efficient.

However, a pipeline can indirectly change the training result if it changes:

• Observation order
• Batch composition
• Dropped observations
• Random augmentation
• Preprocessing
• Numerical data types
• The number of updates per epoch

If accuracy changes, we should investigate which of these experimental conditions changed instead of concluding that tf.data is inherently more accurate.

Evaluating the Test Dataset

After model selection is complete, evaluation works exactly as it did with NumPy arrays:

test_results = model.evaluate(
    test_dataset,
    return_dict=True,
    verbose=0,
)
for metric_name, metric_value in test_results.items():
    print(
        f"{metric_name}: {metric_value:.4f}"
    )

For prediction, a dataset should usually contain only features:

test_feature_dataset = (
    tf.data.Dataset
    .from_tensor_slices(X_test_scaled)
    .batch(BATCH_SIZE)
    .prefetch(tf.data.AUTOTUNE)
)
test_probabilities = model.predict(
    test_feature_dataset,
    verbose=0,
).ravel()

Alternatively, Keras can generate predictions from a dataset containing feature-label pairs, but a feature-only dataset expresses the inference purpose more clearly.

The predicted order matches the dataset iteration order. Because we did not shuffle the test dataset, each prediction can be connected back to its original observation.

Common tf.data Pipeline Mistakes

Shuffling Features and Labels Separately

Features and labels must enter from_tensor_slices() as a shared structure. Independent shuffling destroys their alignment.

Shuffling After Batching

This changes the order of complete batches rather than thoroughly mixing individual observations.

Shuffling Validation or Test Data Without a Reason

It does not improve evaluation and can make row-level error analysis harder.

Using an Inadequate Shuffle Buffer

A very small buffer only mixes nearby observations. This may be insufficient when the source data has a strong ordering.

Using an Excessively Large Buffer

A full-dataset buffer may consume too much memory for a large dataset. Buffer size is a trade-off, not a ritual.

Batching Twice

If the dataset is already batched, do not also pass a batch size to Keras expecting a second normal batching stage.

Caching Data That Does Not Fit

In-memory caching can exhaust memory. File-based caching also requires enough local storage and a deliberate invalidation strategy.

Caching Random Transformations

Caching after a random augmentation or random transformation may freeze one realization and reuse it across epochs. Place the cache according to which operations should be repeated.

Dropping the Final Batch Unintentionally

drop_remainder=True removes observations from each pass whenever the dataset size is not divisible by the batch size.

Calling repeat() Without Defining the Epoch

An indefinitely repeating dataset needs an intentional steps_per_epoch value.

Learning Preprocessing Statistics from All Data

Neither tf.data nor Keras prevents leakage. Imputers, scalers, normalization layers, and vocabularies must learn their state from training data only.

Frequently Asked Questions

What is a tensor in simple terms?

A tensor is a numerical array that may have zero, one, two, or more axes. Structured features are commonly stored as a two-dimensional tensor containing observations and features.

Is tf.data necessary for small datasets?

No. NumPy arrays are often sufficient when data fits in memory and preprocessing is simple. tf.data becomes more valuable when pipelines involve files, repeated transformations, parallel processing, or performance-sensitive training.

What is the best batch size for structured data?

There is no universal best value. Batch size affects gradient noise, update frequency, memory use, and hardware efficiency. It should be compared through controlled validation experiments.

Should validation data be shuffled?

Usually not. Validation does not update the model, and stable ordering simplifies prediction matching and error analysis.

Does prefetching use more memory?

Yes, prefetching maintains a buffer of elements or batches. The additional memory is exchanged for the possibility of overlapping data preparation with model computation.

Should I always use cache and prefetch together?

No. Prefetching is broadly useful when producer and consumer work can overlap. Caching is useful when repeated upstream work is expensive and the cached result fits in memory or storage.

Can tf.data prevent data leakage?

No. It provides pipeline transformations but does not understand experimental boundaries. The practitioner must ensure that learned preprocessing uses training data only.

Why is the final batch smaller?

The dataset size is not always divisible by the batch size. With drop_remainder=False, TensorFlow preserves the remaining observations in a smaller final batch.

Final Perspective

In Part 19, Keras replaced our manually written training loop. In this article, tf.data replaced the assumption that training data must arrive as one ready-made array.

The model still receives feature values and labels. It still calculates predictions, losses, gradients, and parameter updates. What changes is the machinery surrounding that computation.

A well-designed input pipeline makes the movement of data explicit:

• Individual observations become tensors.
• Tensors are shuffled without breaking label alignment.
• Observations are grouped into mini-batches.
• Batches are prepared while earlier batches are being processed.
• Training, validation, and test boundaries remain intact.

The abstraction is useful because we understand what it represents.

shuffle() is not merely a performance command; it changes batch composition. batch() is not merely a memory command; it changes gradient estimation and update frequency. cache() is not automatically beneficial; its value depends on what work occurs before it. prefetch() can improve throughput but does not create a better decision boundary.

In Part 21, we will turn our attention back to the network itself. We will examine Glorot and He initialization, show how initial weight variance affects signals and gradients, and distinguish Batch Normalization from techniques designed primarily to control overfitting.