Introduction

There are two different situations hidden inside this question. Sometimes you are using TensorFlow Datasets, where the dataset already exposes named splits such as train and test. Other times you already have a tf.data.Dataset object and need to split it yourself.

The correct answer depends on which one you have. If predefined splits exist, use them. If you only have one dataset stream, split it deterministically with take and skip, ideally after a seeded shuffle if you need randomness.

If You Are Using TensorFlow Datasets

For datasets loaded through tensorflow_datasets, the cleanest approach is to request the split directly when loading.

1import tensorflow_datasets as tfds
2
3train_ds, test_ds = tfds.load(
4    'mnist',
5    split=['train[:80%]', 'train[80%:]'],
6    as_supervised=True,
7)
8
9for image, label in train_ds.take(1):
10    print(image.shape, label.numpy())

This works well because the split happens at the dataset-definition level instead of after you have already built a pipeline.

If the dataset already provides official train and test splits, prefer those:

train_ds = tfds.load('mnist', split='train', as_supervised=True)
test_ds = tfds.load('mnist', split='test', as_supervised=True)

That is usually better than inventing your own split when the dataset author already defined one.

If You Already Have a tf.data.Dataset

If you created a dataset yourself, you can split it with take and skip.

1import tensorflow as tf
2
3values = tf.data.Dataset.range(10)
4
5train_size = 8
6train_ds = values.take(train_size)
7test_ds = values.skip(train_size)
8
9print(list(train_ds.as_numpy_iterator()))
10print(list(test_ds.as_numpy_iterator()))

This is deterministic and easy to reason about. The first eight items go to training and the rest go to testing.

Add Shuffling Before Splitting When Needed

If the original dataset is ordered by class, time, or source, a raw take and skip split may be biased. In that case, shuffle first with a fixed seed.

1import tensorflow as tf
2
3values = tf.data.Dataset.range(10)
4values = values.shuffle(buffer_size=10, seed=42, reshuffle_each_iteration=False)
5
6train_ds = values.take(8)
7test_ds = values.skip(8)
8
9print(list(train_ds.as_numpy_iterator()))
10print(list(test_ds.as_numpy_iterator()))

The reshuffle_each_iteration=False option matters when you want a stable train-test boundary instead of a different split every epoch.

Batch After the Split

Split first, then batch and preprocess. That keeps the data boundary clear and avoids accidental leakage between train and test.

1import tensorflow as tf
2
3features = tf.data.Dataset.range(100).map(lambda x: (tf.cast(x, tf.float32), x % 2))
4features = features.shuffle(100, seed=123, reshuffle_each_iteration=False)
5
6train_ds = features.take(80).batch(16)
7test_ds = features.skip(80).batch(16)

You can also map preprocessing before batching, but the important part is that the actual split should happen before anything stateful or evaluation-specific is mixed together.

When Cardinality Is Unknown

Some datasets come from generators, streaming sources, or pipelines where the total length is not known up front. In those cases, exact ratio splitting becomes harder.

You then have a few options:

• split earlier, before creating the streaming dataset
• materialize metadata so size is known
• use source-level partitioning such as separate files for train and test

For large or production pipelines, source-level partitioning is usually cleaner than trying to improvise a split late in the tf.data graph.

Why the Order Matters

A train-test split is not just a coding exercise. It protects evaluation integrity. If you split incorrectly after caching, repeating, or reshuffling per epoch, you may accidentally leak examples between training and test.

That gives you optimistic metrics and misleading conclusions about model quality.

Common Pitfalls

A common mistake is calling shuffle() separately on train and test after slicing an ordered dataset. That does not fix a biased split if the boundary itself was already bad.

Another mistake is using reshuffle_each_iteration=True before take and skip when you expect a stable split. That changes the assignment over time.

People also ignore predefined TFDS splits and rebuild them manually. If a dataset already defines train, validation, and test, use them unless you have a strong reason not to.

Finally, be careful with caching and repeating. Those operations can make split bugs harder to notice if placed in the wrong order.

Summary

• For TensorFlow Datasets, prefer predefined splits or percentage-based split= syntax.
• For plain tf.data.Dataset objects, use take and skip.
• Shuffle with a fixed seed before splitting if the original order is biased.
• Split before batching, repeating, or any pipeline logic that could blur dataset boundaries.
• Stable evaluation depends on a stable and leak-free split, not just on getting the code to run.