How to build a TensorFlow dataset input pipeline

Model training can wait on data preparation even when the model itself is small. A tf.data.Dataset keeps feature-label pairs, preprocessing, batching, and prefetching on one iterable that Keras can consume for fitting and validation.

Aligned in-memory tensors work well with Dataset.from_tensor_slices() because it slices every component along the first dimension. The training path can shuffle before batching, while validation keeps its original order; both paths apply the same vectorized scaling after batching.

The starting tensors are already split, and TensorFlow is already installed in the active Python environment. Batched TensorSpec objects with matching label shapes show that the pipeline structure is usable, while loss and val_loss in the fit history show that model.fit() consumed both datasets.

1. Create tf-input-pipeline.py with aligned tensors for both dataset splits.

   tf-input-pipeline.py

   import os
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.keras.utils.set_random_seed(7)
    
   train_features = tf.constant(
       [
           [510.0, 6.5, 1.2],
           [620.0, 7.1, 0.8],
           [470.0, 5.9, 1.7],
           [730.0, 8.0, 0.4],
           [690.0, 7.8, 0.6],
           [540.0, 6.8, 1.0],
           [455.0, 5.7, 1.9],
           [760.0, 8.3, 0.3],
       ],
       dtype=tf.float32,
   )
   train_labels = tf.constant([0, 1, 0, 1, 1, 0, 0, 1], dtype=tf.float32)
    
   validation_features = tf.constant(
       [
           [500.0, 6.4, 1.1],
           [710.0, 7.9, 0.5],
           [480.0, 6.0, 1.6],
           [745.0, 8.1, 0.4],
       ],
       dtype=tf.float32,
   )
   validation_labels = tf.constant([0, 1, 0, 1], dtype=tf.float32)
    
   feature_scale = tf.constant([1000.0, 10.0, 10.0], dtype=tf.float32)

2. Append the batch preprocessing function after feature_scale.

   def prepare_batch(features, labels):
       return features / feature_scale, labels

3. Append the dataset builder below prepare_batch().

   def build_dataset(features, labels, batch_size, training=False):
       dataset = tf.data.Dataset.from_tensor_slices((features, labels))
       if training:
           dataset = dataset.shuffle(
               buffer_size=features.shape[0],
               seed=7,
               reshuffle_each_iteration=True,
           )
       dataset = dataset.batch(batch_size)
       dataset = dataset.map(
           prepare_batch,
           num_parallel_calls=tf.data.AUTOTUNE,
       )
       return dataset.prefetch(tf.data.AUTOTUNE)

   The training flag keeps shuffle() off the validation data while map(), batch(), and prefetch() stay consistent across both splits.
   Related: How to optimize TensorFlow data pipeline performance

4. Append both dataset instances below build_dataset().

   train_dataset = build_dataset(
       train_features,
       train_labels,
       batch_size=4,
       training=True,
   )
   validation_dataset = build_dataset(
       validation_features,
       validation_labels,
       batch_size=2,
   )

5. Append the compiled Keras model below the validation dataset.

   model = tf.keras.Sequential(
       [
           tf.keras.layers.Input(shape=(3,)),
           tf.keras.layers.Dense(4, activation="relu"),
           tf.keras.layers.Dense(1, activation="sigmoid"),
       ]
   )
   model.compile(optimizer="adam", loss="binary_crossentropy")

6. Append the fit-history inspection block below model.compile().

   history = model.fit(
       train_dataset,
       validation_data=validation_dataset,
       epochs=1,
       shuffle=False,
       verbose=0,
   )
    
   train_batch_features, train_batch_labels = next(iter(train_dataset))
   validation_batch_features, validation_batch_labels = next(
       iter(validation_dataset)
   )
   fit_history = {
       key: len(values)
       for key, values in sorted(history.history.items())
   }
    
   print(f"tensorflow={tf.__version__}")
   print(f"train_element_spec={train_dataset.element_spec}")
   print(f"validation_element_spec={validation_dataset.element_spec}")
   print(f"train_batch_shape={tuple(train_batch_features.shape)}")
   print(f"train_label_shape={tuple(train_batch_labels.shape)}")
   print(f"validation_batch_shape={tuple(validation_batch_features.shape)}")
   print(f"validation_label_shape={tuple(validation_batch_labels.shape)}")
   print(f"fit_history={fit_history}")

   validation_data evaluates the validation dataset after the training epoch. The dataset pipeline owns shuffling, so shuffle=False avoids asking Keras to shuffle an already-built dataset.
   Related: How to train, evaluate, and run prediction with a Keras model

7. Run tf-input-pipeline.py to confirm the batched datasets reach model.fit().

   $ python3 tf-input-pipeline.py
   tensorflow=2.21.0
   train_element_spec=(TensorSpec(shape=(None, 3), dtype=tf.float32, name=None), TensorSpec(shape=(None,), dtype=tf.float32, name=None))
   validation_element_spec=(TensorSpec(shape=(None, 3), dtype=tf.float32, name=None), TensorSpec(shape=(None,), dtype=tf.float32, name=None))
   train_batch_shape=(4, 3)
   train_label_shape=(4,)
   validation_batch_shape=(2, 3)
   validation_label_shape=(2,)
   fit_history={'loss': 1, 'val_loss': 1}