keras
neural networks
resource exhausted
model training
GPU memory issues

Resource Exhausted when training a neural network - keras

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In Keras, a ResourceExhausted error usually means the model, batch, or intermediate tensors do not fit into available memory, most often GPU memory. The message often appears during forward or backward passes when TensorFlow tries to allocate a large tensor and cannot. The fix is usually not one magic switch, but a combination of reducing memory pressure and understanding where the memory is going.

What Usually Consumes the Memory

Training memory is not only the model weights. A training step also needs space for:

  • activations from each layer
  • gradients during backpropagation
  • optimizer state such as momentum or Adam statistics
  • the input batch itself

That is why a model that fits for inference can still fail during training. Backpropagation is more memory-intensive than simple forward prediction.

The Fastest Fix: Reduce Batch Size

The first thing to try is usually a smaller batch size.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(128, 128, 3)),
5    tf.keras.layers.Conv2D(32, 3, activation="relu"),
6    tf.keras.layers.Flatten(),
7    tf.keras.layers.Dense(10)
8])
9
10model.compile(optimizer="adam", loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))
11
12# Lower batch size if memory is tight.
13model.fit(x_train, y_train, batch_size=8, epochs=5)

If the error happens at batch_size=64, try 32, 16, or 8. This is often the highest-leverage change because activations scale directly with batch size.

Reduce Model Size or Input Size

If lowering the batch size is not enough, the next step is reducing model footprint.

Common options are:

  • fewer layers
  • fewer filters or hidden units
  • smaller input images
  • replacing a Flatten layer with pooling

For example, this version is often much lighter than a large flatten-based network:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(128, 128, 3)),
5    tf.keras.layers.Conv2D(16, 3, activation="relu"),
6    tf.keras.layers.MaxPooling2D(),
7    tf.keras.layers.Conv2D(32, 3, activation="relu"),
8    tf.keras.layers.GlobalAveragePooling2D(),
9    tf.keras.layers.Dense(10)
10])

A Flatten layer after large feature maps often explodes parameter count and memory usage.

Enable Memory Growth on GPU

TensorFlow can reserve large chunks of GPU memory up front. Allowing memory growth can make notebook and shared-GPU environments behave better.

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4for gpu in gpus:
5    tf.config.experimental.set_memory_growth(gpu, True)

This does not make the GPU larger, but it can reduce wasteful allocation behavior and improve coexistence with other workloads.

Mixed Precision Can Help

On supported hardware, mixed precision can cut memory usage and often speed training too.

python
1import tensorflow as tf
2from tensorflow.keras import mixed_precision
3
4mixed_precision.set_global_policy("mixed_float16")

This is especially useful on modern GPUs designed for lower-precision matrix math. Just remember that output layers or losses may need careful dtype handling in some models.

Clean Up Old Models in Notebook Sessions

If you are iterating in a notebook, old graphs and model objects can accumulate. Clear them before rebuilding a new model.

python
1import gc
2import tensorflow as tf
3
4del model
5tf.keras.backend.clear_session()
6gc.collect()

This is not a substitute for reducing true memory demand, but it helps when experimentation leaves stale state around.

Input Pipelines Matter Too

Sometimes the training input pipeline is the hidden memory problem. Large preloaded arrays, aggressive prefetch settings, or duplicated datasets can consume far more RAM than expected.

TensorFlow datasets are often safer than loading everything into giant Python structures.

python
1train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
2train_ds = train_ds.shuffle(1000).batch(16).prefetch(tf.data.AUTOTUNE)
3
4model.fit(train_ds, epochs=5)

Even here, be mindful that batching still affects device memory usage.

Common Pitfalls

The biggest mistake is focusing only on model weights. During training, activations and optimizer state often consume as much or more memory than the raw parameter tensors.

Another mistake is using a large Flatten layer after high-resolution convolution features. That creates enormous dense layers very quickly.

People also keep retrying the same model in a notebook without clearing old state, which makes memory behavior look worse than it really is.

Finally, do not assume the fix is always a stronger GPU. Often a smaller batch, lighter architecture, or mixed precision solves the immediate issue without changing hardware.

Summary

  • 'ResourceExhausted usually means the training step needs more memory than is available.'
  • Reducing batch size is often the fastest and most effective fix.
  • Smaller models and smaller inputs reduce memory pressure substantially.
  • GPU memory growth and mixed precision can help in the right environments.
  • In notebook workflows, clear old model state before running new experiments.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track 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.

Practice ML system design

All Rights Reserved.