Introduction

Batch size is the number of training examples processed together before one gradient update. In TensorFlow, it affects memory usage, step time, optimization noise, and sometimes model quality, so it is one of the most important training hyperparameters to understand.

What Batch Size Actually Changes

Suppose you have 10,000 training examples and a batch size of 100. One epoch will contain 100 training steps. If the batch size is 1,000, one epoch will contain only 10 steps.

That changes several things at once:

• how much data is seen per optimizer update
• how much GPU or CPU memory is required
• how noisy the gradient estimate is
• how many updates happen per epoch

So batch size is not only a performance setting. It also changes the optimization dynamics.

Set Batch Size with tf.data

In TensorFlow, the most common place to define batch size is the input pipeline:

1import tensorflow as tf
2
3dataset = tf.data.Dataset.from_tensor_slices((
4    tf.random.normal((1000, 20)),
5    tf.random.uniform((1000,), maxval=2, dtype=tf.int32),
6))
7
8dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)

Here, 32 is the batch size. Each training step will receive tensors containing 32 examples.

Batch Size and Memory

Larger batches consume more memory because the framework must store more activations and gradients at once. If you push the batch size too far, you will hit out-of-memory errors before training even starts.

That is why the practical upper bound is often determined by:

• model size
• input resolution
• sequence length
• available GPU RAM

A large image model with a batch size of 128 may be impossible on one GPU, while a small tabular model may handle much larger batches comfortably.

Small Versus Large Batches

There is no universally best batch size. The tradeoff is:

• small batches give noisier gradients and more frequent updates
• large batches give smoother gradients and better hardware throughput

Small batches can sometimes generalize well because the gradient noise acts like a regularizer. Large batches can reduce wall-clock time per epoch, but they may require learning-rate tuning and do not automatically improve final model quality.

Batch Size in model.fit

If you are not using a pre-batched dataset, you can also set batch size directly in Keras training:

1model = tf.keras.Sequential([
2    tf.keras.layers.Dense(32, activation="relu"),
3    tf.keras.layers.Dense(1),
4])
5
6model.compile(optimizer="adam", loss="mse")
7
8x = tf.random.normal((1000, 20))
9y = tf.random.normal((1000, 1))
10
11model.fit(x, y, batch_size=64, epochs=5)

This tells Keras to split the in-memory arrays into batches of 64 for training.

How to Choose a Good Starting Point

A practical workflow is:

1. start with a moderate batch size such as 32 or 64
2. increase it until memory becomes a problem or validation quality drops
3. retune learning rate if batch size changes substantially

If you double the batch size, you should not assume the old optimizer settings are still ideal.

Common Pitfalls

• Treating batch size as only a speed setting and ignoring its effect on optimization.
• Choosing the largest possible batch just because it fits in memory.
• Forgetting that changing batch size often requires retuning the learning rate.
• Comparing experiments by epoch count alone when the number of optimizer updates per epoch has changed.
• Mixing pre-batched datasets with another batch_size setting in training and getting confused about the real input shape.

Summary

• Batch size is the number of samples processed together in one training step.
• It affects memory use, throughput, gradient noise, and convergence behavior.
• In TensorFlow, batch size is usually set with dataset.batch(...) or model.fit(..., batch_size=...).
• Larger is not always better, even when it fits in memory.
• A moderate starting point plus measurement is usually the best way to tune it.