tensorflow
dataset
shuffle
batch processing
machine learning

tensorflow dataset shuffle then batch or batch then shuffle

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In most training pipelines, you should shuffle the dataset before batching. That gives each mini-batch a mixture of examples from across the dataset instead of merely randomizing the order of already-formed batches.

Batching first and then shuffling is usually a weaker form of randomness. It can still move groups around, but the internal contents of each batch stay correlated, which is often not what you want for stochastic training.

Shuffle Then Batch Is the Normal Choice

This is the common TensorFlow pattern:

python
1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(20)
4dataset = dataset.shuffle(buffer_size=20)
5dataset = dataset.batch(4)
6
7for batch in dataset.take(3):
8    print(batch.numpy())

Here, individual examples are shuffled first and then grouped into batches. That means each batch is formed from a randomized stream of items, which is what most SGD-style training expects.

This is usually the right order for:

  • image classification
  • tabular supervised learning
  • language-model mini-batching when full randomization is acceptable
  • most standard training loops

What Happens if You Batch First

Now compare that with batching first:

python
1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(20)
4dataset = dataset.batch(4)
5dataset = dataset.shuffle(buffer_size=5)
6
7for batch in dataset.take(3):
8    print(batch.numpy())

In this version, TensorFlow shuffles batches rather than individual elements. So a batch such as [0, 1, 2, 3] may move later in the order, but those four examples stay together.

That can be useful in special situations, but it does not produce the same training behavior as element-level shuffling.

Why the Difference Matters

If the original dataset is ordered by class, time, user, or file grouping, batching before shuffling can create batches that are internally homogeneous or strongly correlated. That can hurt gradient quality and slow training.

Shuffling before batching breaks those patterns earlier, which usually leads to more representative mini-batches.

The size of the shuffle buffer also matters. A tiny buffer does not produce a fully random permutation, so even the right order can still be weak if the buffer is too small.

Common Training Pipeline Pattern

A typical pipeline is:

python
dataset = dataset.shuffle(buffer_size=10000)
dataset = dataset.batch(32)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

If the dataset repeats for multiple epochs, the order relative to repeat() also matters. A common setup is:

  • 'shuffle() before each epoch'
  • then batch()
  • then prefetch()

The exact best pipeline can depend on data size and whether the data is streaming, but shuffle-before-batch is still the default answer for ordinary supervised training.

When Batch Then Shuffle Can Be Reasonable

There are niche cases where batching first is intentional:

  • examples inside a batch must stay grouped
  • sequence windows were constructed deliberately and should not be broken apart
  • you want to randomize batch order for efficiency but preserve internal structure

That is less common than ordinary example-level shuffling, but it is not always wrong. It is just solving a different problem.

Common Pitfalls

  • Assuming batch().shuffle() is equivalent to shuffle().batch(). It is not.
  • Using a tiny shuffle buffer and expecting near-perfect randomization.
  • Forgetting that correlated examples inside a batch can hurt SGD behavior.
  • Shuffling after batching when the real goal was example-level randomness.
  • Copying a pipeline order from another problem without checking whether the dataset has structure that should or should not be preserved.

Summary

  • In most training pipelines, shuffle before batch.
  • 'shuffle().batch() randomizes examples before groups are formed.'
  • 'batch().shuffle() only randomizes the order of whole batches.'
  • The correct order depends on whether examples should be mixed or preserved in groups.
  • For standard supervised learning, shuffle-then-batch is the usual best default.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.