How to control memory while using Keras with tensorflow backend?
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
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:
If you want a hard memory cap for experiments, create a logical device:
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:
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:
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:
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:
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.datainstead of materializing more than you need. - If memory still fails, simplify the model or use mixed precision on supported hardware.
Related reading
- How to convert a list of tensors of dim N to a tensor of dim N1
- How to convert a PyTorch nn.Module into a HuggingFace PreTrainedModel object?
- How to convert model.module.fc when using DP, if using DDP
- How to convert Tensorflow 2.0 SavedModel to TensorRT?
- How to control tensorflow's VLOG?
- How to control when to compute evaluation vs training using the Estimator API of tensorflow?
- How to convert a double to long without casting?
- How to convert a ListString into a comma separated string without iterating List explicitly

DSA Fundamentals
Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.
View the courseTrack 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.