TensorFlow
Deep Learning
OOM Error
MNIST
Resource Management

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:

  • '10000 examples in one batch'
  • '32 feature maps or channels'
  • '28 x 28 spatial 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.

python
1import tensorflow as tf
2from tensorflow import keras
3
4(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
5x_train = x_train.astype("float32") / 255.0
6x_test = x_test.astype("float32") / 255.0
7
8x_train = x_train[..., None]
9x_test = x_test[..., None]
10
11model = keras.Sequential([
12    keras.layers.Conv2D(32, 3, activation="relu", input_shape=(28, 28, 1)),
13    keras.layers.MaxPooling2D(),
14    keras.layers.Flatten(),
15    keras.layers.Dense(128, activation="relu"),
16    keras.layers.Dense(10, activation="softmax")
17])
18
19model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
20model.fit(x_train, y_train, batch_size=128, epochs=3)

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.

python
1train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(128)
2test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(128)
3
4model.fit(train_ds, epochs=3)
5model.evaluate(test_ds)

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 32 to 16
  • 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.

python
gpus = tf.config.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

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 10000 is 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.

Course illustration
Course illustration

All Rights Reserved.