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

Model evaluation in TensorFlow depends on keeping model fitting, tuning, and final metric reporting on separate examples. A train, validation, and test layout lets Keras learn from one partition, compare model choices against validation data, and reserve the test set for the final metrics run.

For a finite tf.data.Dataset, tf.keras.utils.split_dataset() can divide the dataset directly and return new dataset objects with the same element structure. A 70/15/15 layout is easiest to build as a 70 percent training split followed by a 50/50 split of the remaining holdout data.

Split before batching so Dataset.cardinality() still counts examples instead of batches. Keep a fixed seed when shuffling is enabled, check the resulting counts, and leave published TensorFlow Datasets (TFDS) train, validation, or test partitions intact when the source dataset already defines them.

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

   $ python3 -c "import tensorflow as tf; print(tf.__version__)"
   2.21.0

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

2. Save the dataset split demo as

   tensorflow-dataset-split.py

   .

   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))
    
   train_split, holdout_split = tf.keras.utils.split_dataset(
       full_dataset,
       left_size=0.7,
       shuffle=True,
       seed=7,
   )
   validation_split, test_split = tf.keras.utils.split_dataset(
       holdout_split,
       left_size=0.5,
       shuffle=False,
   )
    
   train_count = int(train_split.cardinality().numpy())
   validation_count = int(validation_split.cardinality().numpy())
   test_count = int(test_split.cardinality().numpy())
    
   train_ds = train_split.batch(4)
   validation_ds = validation_split.batch(4)
   test_ds = test_split.batch(4)
    
   train_features, train_labels = next(iter(train_ds))
   validation_features, validation_labels = next(iter(validation_ds))
   test_features, test_labels = next(iter(test_ds))
    
   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"train_batch_shape={tuple(train_features.shape)}")
   print(f"train_label_shape={tuple(train_labels.shape)}")
   print(f"validation_batch_shape={tuple(validation_features.shape)}")
   print(f"validation_label_shape={tuple(validation_labels.shape)}")
   print(f"test_batch_shape={tuple(test_features.shape)}")
   print(f"test_label_shape={tuple(test_labels.shape)}")

   The first split shuffles the full dataset with seed=7 before creating the training and holdout datasets. The second split divides only the holdout dataset, so the validation and test partitions stay separate from training data.
   Related: How to set a random seed in TensorFlow

3. Run the script and confirm that the split counts add up to the original 20 examples.

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

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

4. Reuse the unbatched splits in the training code before adding project-specific batching and prefetching.

   train_ds = train_split.shuffle(train_count, seed=7).batch(32).prefetch(tf.data.AUTOTUNE)
   validation_ds = validation_split.batch(32).prefetch(tf.data.AUTOTUNE)
   test_ds = test_split.batch(32).prefetch(tf.data.AUTOTUNE)
    
   history = model.fit(train_ds, validation_data=validation_ds, epochs=10)
   test_loss, test_accuracy = model.evaluate(test_ds)

   Pass only the validation split to validation_data during training, and keep the test split isolated until the model architecture and tuning choices are fixed.
   Related: How to train, evaluate, and run prediction with a Keras model
   Related: How to optimize TensorFlow data pipeline performance