Introduction

Keras models can exhaust GPU or host memory for several different reasons: batch size, large activations, repeated graph creation, or TensorFlow reserving most GPU memory up front. Fixing the problem is usually a combination of runtime configuration and model-design changes rather than a single magic setting.

Control GPU Allocation Early

TensorFlow can be configured to grow GPU memory usage as needed instead of reserving everything immediately. This must happen before TensorFlow initializes the GPU:

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)
6
7print(gpus)

If you want a hard memory cap for experiments, create a logical device:

1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4if gpus:
5    tf.config.set_logical_device_configuration(
6        gpus[0],
7        [tf.config.LogicalDeviceConfiguration(memory_limit=4096)]
8    )

That sets a four-gigabyte limit on the first GPU. It is useful when running multiple jobs on one machine or when you want predictable failure boundaries during debugging.

Reduce the Biggest Driver: Batch Size

Batch size is usually the fastest lever to pull:

1history = model.fit(
2    train_ds,
3    epochs=5,
4    batch_size=16
5)

Smaller batches reduce the memory held for activations during forward and backward passes. If training becomes unstable, you can often compensate with gradient accumulation or a learning-rate adjustment instead of immediately returning to the original batch size.

Avoid Building Graphs Repeatedly

Interactive notebooks and hyperparameter loops often leak memory because new models are created over and over without clearing the previous state:

1import gc
2import tensorflow as tf
3from tensorflow import keras
4
5for units in [32, 64, 128]:
6    keras.backend.clear_session()
7
8    model = keras.Sequential([
9        keras.layers.Input(shape=(100,)),
10        keras.layers.Dense(units, activation="relu"),
11        keras.layers.Dense(1)
12    ])
13
14    model.compile(optimizer="adam", loss="mse")
15    model.fit(x_train, y_train, epochs=1, batch_size=16, verbose=0)
16
17    del model
18    gc.collect()

clear_session() is especially important when you are repeatedly creating models inside the same Python process.

Stream Data Instead of Materializing Everything

Large NumPy arrays or image datasets can consume memory before training even starts. Use tf.data to stream and prefetch:

1import tensorflow as tf
2
3train_ds = (
4    tf.data.Dataset.from_tensor_slices((x_train, y_train))
5    .shuffle(10000)
6    .batch(16)
7    .prefetch(tf.data.AUTOTUNE)
8)

This does not make the model itself smaller, but it prevents the input pipeline from becoming a second independent memory problem.

Make the Model Cheaper

Some models are simply too large for the available hardware. Common changes that reduce memory pressure include:

• smaller input resolution
• fewer filters or hidden units
• fewer layers
• mixed precision on supported GPUs

Mixed precision can be enabled like this:

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

This can reduce memory usage and sometimes improve performance, but you should test training stability and final accuracy.

Distinguish Training Memory from Inference Memory

Training uses much more memory than inference because intermediate activations must be stored for backpropagation. If a model fits for prediction but fails during fit, the likely culprits are:

• batch size
• optimizer state
• backpropagation activations

That distinction helps avoid chasing the wrong fixes, such as compressing input files when the real issue is model state during gradient computation.

Common Pitfalls

• Calling the GPU memory-growth configuration after TensorFlow has already initialized the device.
• Rebuilding many models in one process without clear_session().
• Loading entire datasets into RAM when a streaming pipeline would work.
• Assuming the input data is the only issue when activations and optimizer state dominate memory use.
• Enabling mixed precision without verifying numerical stability or metric behavior.

Summary

• Configure TensorFlow GPU memory behavior before the first GPU operation.
• Reduce batch size first because it often gives the quickest memory relief.
• Clear Keras session state when creating models repeatedly in the same process.
• Stream data with tf.data instead of materializing more than you need.
• If memory still fails, simplify the model or use mixed precision on supported hardware.