Keras
TensorFlow
Sequential Models
Memory Management
Machine Learning

Keras TensorFlow, CPU Training Sequential models in loop eats memory

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

Training Keras Sequential models in a loop can make memory usage climb when each iteration creates new models, graphs, histories, datasets, or traced functions without releasing them. On CPU-only systems this is especially noticeable in long experiments, where the issue is often not one dramatic leak but steady accumulation of TensorFlow state and Python references.

Why Memory Grows in Training Loops

Typical causes include:

  • creating a brand-new model every iteration
  • compiling repeatedly
  • keeping references to history objects, callbacks, or tensors
  • rebuilding datasets over and over
  • triggering retracing because shapes or signatures keep changing

Even if each individual object is small, long runs can accumulate enough state to make the process look leaky.

Clear Backend State Between Independent Runs

If each loop iteration is meant to be an independent training run, clear TensorFlow state and release references:

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

This is not a cure for every possible TensorFlow memory issue, but it removes several common sources of accumulation.

Reuse Models When the Experiment Allows It

If the architecture stays the same and only the input data changes, reusing one model is often cheaper than rebuilding it repeatedly:

python
1from tensorflow import keras
2
3model = keras.Sequential([
4    keras.layers.Input(shape=(20,)),
5    keras.layers.Dense(32, activation="relu"),
6    keras.layers.Dense(1)
7])
8
9model.compile(optimizer="adam", loss="mse")
10
11for chunk_x, chunk_y in data_splits:
12    model.fit(chunk_x, chunk_y, epochs=1, verbose=0)

This avoids repeated graph construction and compile overhead.

Watch Out for Lingering References

Sometimes the model is not the only thing being kept alive. Memory may continue growing because code keeps appending:

  • 'History objects'
  • callbacks with stored state
  • prediction tensors
  • large numpy arrays used for evaluation

If you do not need them after each iteration, release them explicitly instead of storing them in a long-lived list.

Stable Shapes Reduce Retracing

TensorFlow may retrace functions when input signatures keep changing. That can increase memory use and slow execution. Keeping input shapes consistent across iterations helps reduce that churn.

This is especially relevant when:

  • batch sizes vary unpredictably
  • custom functions are defined inside loops
  • the same training path is rebuilt with slightly different signatures each time

Build tf.data Pipelines Deliberately

Repeatedly rebuilding datasets can also contribute to memory growth. If the same pipeline can be reused, do so:

python
1import tensorflow as tf
2
3train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)) \
4    .batch(128) \
5    .prefetch(tf.data.AUTOTUNE)
6
7for _ in range(5):
8    model.fit(train_ds, epochs=1, verbose=0)

This usually behaves better than rebuilding a fresh dataset object every loop iteration without reason.

Process Isolation Is the Strongest Reset

For hyperparameter sweeps or many independent model runs, the most reliable reset is often process isolation. Run each experiment in its own process so the operating system reclaims everything when that process exits:

python
1import multiprocessing as mp
2
3def run_one(config):
4    # build, train, evaluate, save results, then exit
5    pass
6
7if __name__ == "__main__":
8    configs = [{"lr": 1e-3}, {"lr": 1e-4}]
9    with mp.Pool(processes=2) as pool:
10        pool.map(run_one, configs)

This is often simpler and more reliable than trying to make one long-lived process perfectly clean across many unrelated experiments.

Measure Memory While You Debug

Use process-level memory measurements so you can distinguish real growth from guesses:

python
1import os
2import psutil
3
4proc = psutil.Process(os.getpid())
5print("rss_mb", proc.memory_info().rss / (1024 * 1024))

Watching memory at the end of each loop iteration often reveals whether cleanup is helping.

Common Pitfalls

  • Rebuilding and recompiling a new model every loop iteration without clearing backend state.
  • Keeping history objects, callbacks, or tensors in long-lived collections.
  • Recreating dataset pipelines unnecessarily.
  • Allowing input shapes to vary in ways that trigger repeated retracing.
  • Assuming Python garbage collection alone will clean up all TensorFlow runtime state.

Summary

  • Memory growth in looped Keras training usually comes from accumulated state, not only one obvious leak.
  • Use tf.keras.backend.clear_session() when runs are independent.
  • Reuse models and datasets when the experiment design allows it.
  • Keep input signatures stable to reduce retracing overhead.
  • For large experiment sweeps, separate processes are often the most reliable cleanup strategy.

Related reading
Course
Intermediate
27 lessons
15 hours
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 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.