Tensorflow Deep MNIST Resource exhausted OOM when allocating tensor with shape10000,32,28,28
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
A TensorFlow ResourceExhaustedError with a shape like [10000, 32, 28, 28] is an out-of-memory problem, not a model-logic problem. The immediate fix is usually to lower the batch size or reduce the size of intermediate tensors, because TensorFlow's own documentation recommends smaller batches or smaller model dimensions when an OOM error occurs.
Read The Tensor Shape Correctly
The shape [10000, 32, 28, 28] usually means:
- '
10000examples in one batch' - '
32feature maps or channels' - '
28 x 28spatial dimensions'
That is already a large activation tensor, and it is only one part of the model's memory usage. Training also needs gradients, optimizer state, temporary buffers, and parameter storage.
If the tensor uses float32, this one allocation alone is roughly 10000 * 32 * 28 * 28 * 4 bytes, which is about 100 MB before the rest of the graph is counted.
The Most Common Cause: Batch Size Is Too Large
A batch dimension of 10000 is the first red flag. For MNIST, that often means code tried to push the full test set or a huge training chunk through the convolution stack at once.
Use mini-batches instead.
A batch size such as 64 or 128 is much more typical.
Use A tf.data Pipeline
Instead of materializing giant batches manually, let TensorFlow stream data in smaller chunks.
This is cleaner and reduces the chance of accidentally building oversized tensors.
Reduce Model Memory If Needed
If smaller batches are not enough, simplify the model.
Examples:
- reduce channel counts such as
32to16 - reduce dense layer width
- remove unnecessary layers
- use mixed precision on supported hardware
For MNIST, you usually do not need a large network.
GPU-Specific Considerations
On GPUs, TensorFlow can also appear to run out of memory because the device is already full. In that case, reduce batch size further or configure memory growth.
This does not create more memory, but it can make allocation behavior less aggressive.
Watch Out For Evaluation Code Too
Training code often uses batches correctly while evaluation or debugging code does not. For example, calling model.predict(x_test) with a very large array can still create oversized intermediate activations if you bypass dataset batching.
Using model.predict(x_test, batch_size=128) or a batched dataset keeps inference memory under control as well.
Common Pitfalls
A common mistake is evaluating or predicting on the entire test set as one giant batch. Even if the dataset is small conceptually, intermediate convolution tensors can still be large.
Another mistake is focusing only on model weights. Activations and gradients often dominate memory during training.
It is also easy to confuse an OOM with a bug in TensorFlow itself. Most of the time, the model or batching strategy simply exceeds available memory.
Summary
- The shape in the error message tells you which tensor is too large.
- A batch size of
10000is usually the main problem in this MNIST-style OOM. - Reduce batch size first, then consider shrinking the model.
- Use
tf.data.Dataset.batch()to avoid accidental giant batches. - On GPUs, memory growth and smaller batches often resolve the issue.

