How to build a TensorFlow dataset input pipeline

TensorFlow training code needs prepared examples to reach the model with a consistent feature shape and label pairing. A tf.data.Dataset input pipeline turns already-split arrays into repeatable batches, with preprocessing and buffering attached to the same data path that Keras consumes during fitting and validation.

Dataset.from_tensor_slices() is the source step for small in-memory arrays because it slices aligned feature and label tensors along their first dimension. After that, map() handles TensorFlow-native preprocessing, shuffle() belongs only on the training split, batch() sets the shape consumed by the model, and prefetch(tf.data.AUTOTUNE) lets TensorFlow overlap input work with training.

The in-memory path assumes the data has already been split into training and validation arrays and still fits in memory. When the source is a CSV file, image directory, or record store, keep the same downstream order after the reader step and verify the result by checking dataset specs, batch shapes, and a short model.fit() run.

1. Open a terminal in a Python environment where TensorFlow already imports cleanly.

   $ python3 - <<'PY'
   import tensorflow as tf
   print(tf.__version__)
   PY
   2.21.0

   Related: How to create a virtual environment for TensorFlow
   Related: How to install TensorFlow with pip

2. Save the pipeline demo as

   tf-input-pipeline.py

   .

   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)
    
    
   def prepare_example(features, label):
       features = tf.cast(features, tf.float32) / feature_scale
       label = tf.cast(label, tf.float32)
       return features, label
    
    
   def build_dataset(features, labels, batch_size, training):
       dataset = tf.data.Dataset.from_tensor_slices((features, labels))
       dataset = dataset.map(prepare_example, num_parallel_calls=tf.data.AUTOTUNE)
       if training:
           dataset = dataset.shuffle(
               buffer_size=int(features.shape[0]),
               seed=7,
               reshuffle_each_iteration=True,
           )
       dataset = dataset.batch(batch_size)
       dataset = dataset.prefetch(tf.data.AUTOTUNE)
       return 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,
       training=False,
   )
    
   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")
    
   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))
   train_feature_spec, train_label_spec = train_dataset.element_spec
   validation_feature_spec, validation_label_spec = validation_dataset.element_spec
    
   print(f"tensorflow={tf.__version__}")
   print(
       "train_feature_spec="
       f"shape={train_feature_spec.shape}, dtype={train_feature_spec.dtype.name}"
   )
   print(
       "train_label_spec="
       f"shape={train_label_spec.shape}, dtype={train_label_spec.dtype.name}"
   )
   print(
       "validation_feature_spec="
       f"shape={validation_feature_spec.shape}, dtype={validation_feature_spec.dtype.name}"
   )
   print(
       "validation_label_spec="
       f"shape={validation_label_spec.shape}, dtype={validation_label_spec.dtype.name}"
   )
   print(f"train_batches={int(train_dataset.cardinality())}")
   print(f"validation_batches={int(validation_dataset.cardinality())}")
   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"history_keys={sorted(history.history.keys())}")
   print(f"fit_epochs={len(history.epoch)}")

   The training flag keeps shuffle() off the validation dataset while still applying the same preprocessing, batching, and prefetching path to both splits.
   Related: How to optimize TensorFlow data pipeline performance

3. Run the script and confirm that both datasets report batched specs before Keras trains for one epoch.

   $ python3 tf-input-pipeline.py
   tensorflow=2.21.0
   train_feature_spec=shape=(None, 3), dtype=float32
   train_label_spec=shape=(None,), dtype=float32
   validation_feature_spec=shape=(None, 3), dtype=float32
   validation_label_spec=shape=(None,), dtype=float32
   train_batches=2
   validation_batches=2
   train_batch_shape=(4, 3)
   train_label_shape=(4,)
   validation_batch_shape=(2, 3)
   validation_label_shape=(2,)
   history_keys=['loss', 'val_loss']
   fit_epochs=1

   The None dimension in each spec line is the variable batch axis. The history_keys and fit_epochs lines confirm that model.fit() consumed the training dataset and evaluated the validation dataset.

4. Use the helper with real arrays and pass the dataset objects directly to model.fit().

   train_ds = build_dataset(
       train_features,
       train_labels,
       batch_size=32,
       training=True,
   )
   validation_ds = build_dataset(
       validation_features,
       validation_labels,
       batch_size=32,
       training=False,
   )
    
   history = model.fit(
       train_ds,
       validation_data=validation_ds,
       epochs=10,
       shuffle=False,
   )

   validation_split is not the right boundary for already-built tf.data.Dataset inputs; keep validation as its own dataset and pass it through validation_data=validation_ds.
   Related: How to train, evaluate, and run prediction with a Keras model
   Related: How to export a SavedModel in TensorFlow

   Keep the same feature order, scaling rule, and label meaning across every split or the pipeline can stay technically valid while the model trains on mismatched inputs.

5. Remove the demo script when it was created only to verify the pipeline shape.

   $ rm tf-input-pipeline.py