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.
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:
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:
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:
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 toshuffle().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
- Tensorflow Dataset.from_generator fails with pyfunc exception
- Tensorflow DecodeJpeg method gives different pixel values on desktop and mobile for the same image
- Tensorflow Deep MNIST Resource exhausted OOM when allocating tensor with shape10000,32,28,28
- tensorflow deep neural network for regression always predict same results in one batch
- Tensorflow dense gradient explanation?
- Tensorflow dense_to_sparse
- Tensorflow device CUDA0 not supported by XLA service while setting up XLA_GPU_JIT device number 0
- Tensorflow Diagonal Subtensor for 3D Convolutional NN
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.