tensorflow
tensorflow-datasets
train-validation-split
machine-learning
tf2.1

Split train data to train and validation by using tensorflow_datasets.load TF 2.1

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

tensorflow_datasets.load can split a dataset for you without manual slicing code. The key is the split argument, which supports percentage-based slicing so you can load a training subset and a validation subset directly from the same source split.

Use Split Expressions Directly

If a dataset has a train split, you can divide it into two pieces with percentage syntax:

python
1import tensorflow_datasets as tfds
2
3train_ds, val_ds = tfds.load(
4    "mnist",
5    split=["train[:80%]", "train[80%:]"],
6    as_supervised=True
7)
8
9print(train_ds)
10print(val_ds)

This is the most common solution. It keeps the split logic inside tfds.load instead of loading the whole training set and partitioning it later.

Add Preprocessing and Batching

After loading the datasets, build the usual input pipeline:

python
1import tensorflow as tf
2import tensorflow_datasets as tfds
3
4
5def preprocess(image, label):
6    image = tf.cast(image, tf.float32) / 255.0
7    return image, label
8
9
10train_ds, val_ds = tfds.load(
11    "mnist",
12    split=["train[:80%]", "train[80%:]"],
13    as_supervised=True
14)
15
16train_ds = train_ds.map(preprocess).shuffle(1000).batch(64).prefetch(tf.data.AUTOTUNE)
17val_ds = val_ds.map(preprocess).batch(64).prefetch(tf.data.AUTOTUNE)

That keeps training shuffled while validation stays stable and deterministic.

When a Dataset Already Has Validation or Test Splits

Some TensorFlow Datasets already provide predefined splits such as train, validation, or test. If that exists, use it instead of inventing a new validation cut from the training set.

For example, a dataset may support:

python
1train_ds, val_ds = tfds.load(
2    "imdb_reviews",
3    split=["train", "test"],
4    as_supervised=True
5)

The important point is to inspect the dataset metadata first rather than assuming you always need percentage slicing.

Splitting More Than Two Ways

You can also create train, validation, and test slices from one original split:

python
splits = ["train[:70%]", "train[70%:85%]", "train[85%:]"]
train_ds, val_ds, test_ds = tfds.load("mnist", split=splits, as_supervised=True)

This is handy for experiments where the dataset ships with only one large training split.

Inspect Dataset Metadata First

Before choosing split expressions, inspect what TFDS already exposes:

python
1import tensorflow_datasets as tfds
2
3builder = tfds.builder("mnist")
4builder.download_and_prepare()
5print(builder.info.splits)

That helps you answer two practical questions. First, does the dataset already provide validation or test. Second, is the dataset large enough that an 80/20 split makes sense for your experiment. Reading the metadata first prevents accidental assumptions about available splits and example counts.

It also makes code reviews easier because anyone reading the training script can see whether the split was chosen intentionally or copied from another project without checking the dataset.

Why This Is Better Than Manual Counting

Using TFDS split expressions makes the pipeline declarative. You do not need to count examples manually, and the split stays close to the data-loading call instead of being scattered across the notebook or training script.

That also makes experiments easier to reproduce because the split logic is visible in one place.

Common Pitfalls

  • Forgetting that training data should usually be shuffled after loading, while validation data usually should not be.
  • Creating a custom validation split even when the dataset already provides an official validation or test split.
  • Mixing percentage-based slicing with assumptions about exact example counts.
  • Applying heavy random augmentation to validation data and then treating validation metrics as stable.
  • Forgetting as_supervised=True when the rest of the pipeline expects (features, label) pairs.

Summary

  • 'tfds.load supports percentage-based split expressions such as train[:80%] and train[80%:].'
  • This is the simplest way to create train and validation datasets from one source split.
  • Build the rest of the input pipeline with map, shuffle, batch, and prefetch.
  • Prefer predefined dataset validation or test splits when they already exist.
  • Keep the split definition close to the load call so experiments stay reproducible.

Course illustration
Course illustration

All Rights Reserved.