TensorFlow Batching and Shuffling
Batching and shuffling are two pipeline operations that control how training data reaches the model. The batch size determines how many samples the model sees before updating its weights. Shuffling prevents the model from learning the order of examples rather than the patterns within them. Both decisions affect training speed, convergence stability, and final model accuracy.
Why Process Data in Batches
Training on one sample at a time (batch size 1) updates weights extremely frequently but with very noisy gradient estimates. Training on the entire dataset at once (full batch) gives a perfect gradient but requires all data to fit in GPU memory and makes only one weight update per epoch. Batches strike the balance — enough samples for a stable gradient, small enough to fit in memory and update weights frequently.
Diagram — Effect of Batch Size on Gradient Estimates: Full Batch (all 10,000 samples): Gradient direction: ──────────────────► exact direction to minimum Updates per epoch: 1 Memory needed: entire dataset Mini-batch (batch=32): Gradient direction: ──~──~────~───────► noisy but roughly correct Updates per epoch: 313 Memory needed: 32 samples at a time Stochastic (batch=1): Gradient direction: ~─~~─~─~~─~~~~~~~► very noisy Updates per epoch: 10,000 Memory needed: 1 sample
Batching in tf.data
import tensorflow as tf
import numpy as np
x = np.random.random((1000, 20)).astype('float32')
y = np.random.randint(0, 3, size=1000)
dataset = tf.data.Dataset.from_tensor_slices((x, y))
# Basic batch
batched = dataset.batch(32)
# Inspect the first batch
for features, labels in batched.take(1):
print(features.shape) # (32, 20)
print(labels.shape) # (32,)
drop_remainder Parameter
# 1000 samples ÷ 32 = 31 full batches + 1 partial batch of 8 # drop_remainder=True drops the partial batch dataset.batch(32, drop_remainder=True) # → Always produces batches of exactly 32 # → Useful when model architecture requires a fixed batch size # drop_remainder=False (default) keeps the partial batch dataset.batch(32, drop_remainder=False) # → Last batch may be smaller
Shuffling in tf.data
# buffer_size controls how many elements are loaded into memory for shuffling # A larger buffer gives better randomness but uses more memory # Perfect shuffle (if data fits in memory) dataset = dataset.shuffle(buffer_size=len(x)) # buffer = full dataset # Approximate shuffle (memory-limited) dataset = dataset.shuffle(buffer_size=5000) # Hold 5000, draw randomly # Reproducible shuffle (for debugging) dataset = dataset.shuffle(buffer_size=1000, seed=42)
How the Shuffle Buffer Works
Dataset (ordered): [1][2][3][4][5][6][7][8][9][10]... buffer_size = 4: fills buffer with first 4 elements Buffer: [1][2][3][4] Draw randomly → outputs [3] Fills slot: [1][2][5][4] Draw randomly → outputs [1] Fills slot: [6][2][5][4] ...and so on Small buffer = elements only shuffled with nearby elements Large buffer = elements shuffled across the whole dataset
The Correct Order: Shuffle → Batch → Prefetch
# CORRECT order
dataset = (
tf.data.Dataset.from_tensor_slices((x, y))
.shuffle(buffer_size=1000) # Shuffle individual examples first
.batch(32) # Then group into batches
.prefetch(tf.data.AUTOTUNE) # Prepare next batch in background
)
# WRONG order — shuffles entire batches, not individual examples
dataset = (
tf.data.Dataset.from_tensor_slices((x, y))
.batch(32)
.shuffle(100) # Shuffles 100 pre-formed batches — much weaker
)
Choosing the Right Batch Size
Batch Size Behavior Typical Use ──────────────────────────────────────────────────────────────── 1 Most noisy gradients, slowest Rarely used 8–16 High noise, many updates Very small datasets 32 Good default balance Most tasks (start here) 64 Stable, fewer updates Medium datasets, large models 128–256 Very stable, fewer updates Large datasets, fast GPUs 512+ Needs large learning rate adj. Distributed training, TPUs ────────────────────────────────────────────────────────────────
Learning Rate Scaling with Batch Size
When you double the batch size, the gradient estimate becomes twice as stable — essentially you get double the signal per step. To compensate, the learning rate should also scale up proportionally. This linear scaling rule works well in practice.
# Linear scaling rule base_batch_size = 32 base_lr = 0.001 new_batch_size = 128 new_lr = base_lr * (new_batch_size / base_batch_size) # new_lr = 0.001 × 4 = 0.004 optimizer = tf.keras.optimizers.Adam(learning_rate=new_lr)
Re-Shuffling Every Epoch
When you use model.fit() with epochs > 1 on a tf.data.Dataset, TensorFlow automatically reshuffles the data at the start of each epoch — as long as reshuffle_each_iteration=True (the default).
# This re-shuffles at every epoch start (default behavior)
dataset = tf.data.Dataset.from_tensor_slices((x, y)).shuffle(1000)
# To disable re-shuffling (rarely needed)
dataset = tf.data.Dataset.from_tensor_slices((x, y)).shuffle(
1000, reshuffle_each_iteration=False
)
Gradient Accumulation: Simulating Large Batches
If you want the stability of a large batch but cannot fit it in GPU memory, gradient accumulation simulates a large effective batch by summing gradients from several small batches before applying one update.
accumulation_steps = 4 # Effective batch = 32 × 4 = 128
for step, (x_batch, y_batch) in enumerate(train_ds):
with tf.GradientTape() as tape:
preds = model(x_batch, training=True)
loss = loss_fn(y_batch, preds) / accumulation_steps # Scale loss
grads = tape.gradient(loss, model.trainable_variables)
# Accumulate gradients
if step == 0:
accumulated = grads
else:
accumulated = [a + g for a, g in zip(accumulated, grads)]
# Apply update every N steps
if (step + 1) % accumulation_steps == 0:
optimizer.apply_gradients(zip(accumulated, model.trainable_variables))
accumulated = None
Batching Sequences of Variable Length
# Sequences can have different lengths — padding aligns them in a batch
sequences = [
[1, 2, 3],
[4, 5],
[6, 7, 8, 9]
]
# padded_batch pads shorter sequences to the length of the longest in the batch
seq_ds = tf.data.Dataset.from_generator(
lambda: iter(sequences),
output_signature=tf.RaggedTensorSpec(shape=[None], dtype=tf.int32)
)
padded = seq_ds.padded_batch(batch_size=3, padded_shapes=[None])
for batch in padded:
print(batch)
# [[1, 2, 3, 0],
# [4, 5, 0, 0],
# [6, 7, 8, 9]]
Getting batching and shuffling right is foundational to stable training. A model that sees data in the same order every epoch learns the sequence rather than the content. A batch size that is too large produces overly smooth gradients that miss fine-grained patterns. These two simple settings have a larger effect on training quality than most architectural choices. The next topic begins the CNN section by covering Conv2D layers in detail.