How to split a dataset into train, validation, and test sets in TensorFlow

A trustworthy model evaluation needs examples that played no part in fitting or tuning. Training data adjusts model weights, validation data guides model choices, and test data measures the chosen model only after those choices are fixed.

For a finite tf.data.Dataset, tf.keras.utils.split_dataset() creates two partitions at a time while preserving each element's feature-and-label structure. A 70/15/15 layout comes from a 70 percent training split followed by an even division of the remaining holdout examples.

Split before batching so Dataset.cardinality() counts examples rather than batches. Keep source-defined TensorFlow Datasets train, validation, and test partitions intact instead of recombining them when those official splits already exist.

1. Create tensorflow-dataset-split.py with a finite dataset containing one unique feature ID per example.

   tensorflow-dataset-split.py

   import tensorflow as tf
    
   features = tf.reshape(tf.cast(tf.range(40), tf.float32), (20, 2))
   labels = tf.cast(tf.range(20) % 2, tf.int32)
   full_dataset = tf.data.Dataset.from_tensor_slices((features, labels))

   The generated tensors stand in for the project's unbatched finite dataset. A stable unique field per example supports the final overlap check.
   Related: How to create a dataset from tensors in TensorFlow

2. Add the training and holdout split below the full_dataset assignment.

   train_split, holdout_split = tf.keras.utils.split_dataset(
       full_dataset,
       left_size=0.7,
       shuffle=True,
       seed=7,
   )

   A fixed seed reproduces the shuffled assignment when the input order and dataset contents remain unchanged.
   Related: How to set a random seed in TensorFlow

3. Divide the holdout split below the first split_dataset() call.

   validation_split, test_split = tf.keras.utils.split_dataset(
       holdout_split,
       left_size=0.5,
       shuffle=False,
   )

   The first call leaves 30 percent in holdout_split, and the second call divides that remainder evenly between validation and test data.

4. Confirm the source dataset has known finite cardinality below the holdout split.

   full_count = int(full_dataset.cardinality())
   if full_count < 0:
       raise ValueError("full_dataset must have finite, known cardinality")

   Dataset.cardinality() can return an unknown or infinite sentinel for some generated, filtered, repeated, or file-backed pipelines. The guard stops those sentinel values from being treated as example counts.

5. Record the three partition cardinalities below the finite-cardinality guard.

   train_count = int(train_split.cardinality())
   validation_count = int(validation_split.cardinality())
   test_count = int(test_split.cardinality())

6. Check partition membership below the cardinality block.

   train_ids = {int(row[0]) for row, _ in train_split}
   validation_ids = {int(row[0]) for row, _ in validation_split}
   test_ids = {int(row[0]) for row, _ in test_split}
   all_ids = train_ids | validation_ids | test_ids
    
   partitions_disjoint = (
       train_ids.isdisjoint(validation_ids)
       and train_ids.isdisjoint(test_ids)
       and validation_ids.isdisjoint(test_ids)
   )
   coverage_complete = len(all_ids) == full_count

   row[0] represents the unique example identifier and may use a different field in project data. Repeated class labels are not suitable membership IDs because many examples can share the same label.

7. Batch each partition below the membership checks.

   train_ds = train_split.batch(4).prefetch(tf.data.AUTOTUNE)
   validation_ds = validation_split.batch(4).prefetch(tf.data.AUTOTUNE)
   test_ds = test_split.batch(4).prefetch(tf.data.AUTOTUNE)

   train_ds supplies model fitting, validation_ds supplies validation_data during tuning, and test_ds remains isolated until final evaluation.
   Related: How to train, evaluate, and run prediction with a Keras model
   Related: How to optimize TensorFlow data pipeline performance

8. Append the split verification report to the end of tensorflow-dataset-split.py.

   print(f"train_examples={train_count}")
   print(f"validation_examples={validation_count}")
   print(f"test_examples={test_count}")
   print(f"total_examples={train_count + validation_count + test_count}")
   print(f"partitions_disjoint={partitions_disjoint}")
   print(f"coverage_complete={coverage_complete}")
   print(f"train_batch_shape={tuple(next(iter(train_ds))[0].shape)}")
   print(f"validation_batch_shape={tuple(next(iter(validation_ds))[0].shape)}")
   print(f"test_batch_shape={tuple(next(iter(test_ds))[0].shape)}")

9. Run tensorflow-dataset-split.py to verify the 70/15/15 counts, disjoint membership, coverage of all 20 source examples, and batch shapes.

   $ python3 tensorflow-dataset-split.py
   train_examples=14
   validation_examples=3
   test_examples=3
   total_examples=20
   partitions_disjoint=True
   coverage_complete=True
   train_batch_shape=(4, 2)
   validation_batch_shape=(3, 2)
   test_batch_shape=(3, 2)

   Fractional split sizes become whole-example counts, so small datasets may not preserve the requested percentages exactly.