TensorFlow
memory leak
session management
resource management
troubleshooting

Tensorflow Memory leak even while closing Session?

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

If TensorFlow memory keeps growing even after sess.close(), the problem is often not a true low-level leak in the session object itself. More commonly, the process is still holding references to graphs, tensors, datasets, or models, or the TensorFlow allocator is intentionally caching memory for reuse. Closing the session is only one part of cleanup, especially in TensorFlow 1.x style code.

Why Closing the Session Is Not Always Enough

A TensorFlow Session owns runtime resources for executing a graph, but it does not automatically erase every Python object connected to that graph. Memory can remain high for several reasons:

  • the default graph keeps growing in a loop
  • Python references still point to tensors or ops
  • Keras models accumulate between runs
  • the GPU allocator keeps memory reserved for reuse
  • input pipelines or generators remain alive

So if you repeatedly build models or graphs without resetting the surrounding state, sess.close() alone will not stop memory growth.

Classic Graph-Growth Problem

One of the most common mistakes in TensorFlow 1.x is building new graph nodes inside a loop:

python
1import tensorflow as tf
2
3for _ in range(1000):
4    x = tf.constant(1.0)
5    y = x + 1.0

This adds new nodes to the default graph every iteration. Even if sessions are opened and closed later, the graph object itself keeps getting larger. A safer pattern is to build the graph once or create a separate graph explicitly:

python
1import tensorflow as tf
2
3graph = tf.Graph()
4with graph.as_default():
5    x = tf.constant(1.0)
6    y = x + 1.0
7
8with tf.compat.v1.Session(graph=graph) as sess:
9    print(sess.run(y))

That keeps graph ownership explicit and limits accidental accumulation.

Resetting State Between Runs

If you are repeatedly creating models in a notebook, training loop, or hyperparameter search, clear framework state as well:

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

clear_session() is especially important when using Keras on top of TensorFlow because Keras keeps global layer and graph state. Without clearing it, each new model can leave behind references that prevent cleanup.

For TensorFlow 1.x style code, you may also need:

python
tf.compat.v1.reset_default_graph()

Use that only when you understand that you are discarding the current default graph. It is a reset tool, not something to sprinkle everywhere blindly.

GPU Memory Behavior Is Often Misread as a Leak

On GPUs, TensorFlow may reserve a large memory pool and keep it instead of immediately returning it to the operating system. That can look like a leak in system monitors even when TensorFlow is simply caching memory for later operations.

You can reduce that behavior by enabling memory growth:

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 solve every memory problem, but it often makes GPU usage easier to reason about during iterative development.

Watch for Hidden Python References

Python objects can keep tensors alive longer than expected. Lists of losses, callbacks that retain outputs, closures over tensors, and global variables are all common sources of accidental retention.

For example:

python
1saved_outputs = []
2
3for _ in range(10):
4    tensor = tf.constant([1.0, 2.0, 3.0])
5    saved_outputs.append(tensor)

Even if sessions are closed elsewhere, the list still holds those tensor objects. If you are profiling a memory issue, inspect the surrounding Python code as seriously as the TensorFlow API usage.

Practical Cleanup Pattern

For repeated training jobs in one process, a safer cleanup sequence looks like this:

python
1import gc
2import tensorflow as tf
3
4
5def train_once():
6    model = tf.keras.Sequential([
7        tf.keras.layers.Dense(32, activation="relu"),
8        tf.keras.layers.Dense(1),
9    ])
10    model.compile(optimizer="adam", loss="mse")
11    return model
12
13
14model = train_once()
15del model
16tf.keras.backend.clear_session()
17gc.collect()

This is not magic, but it covers the most common retained-state paths in practical TensorFlow code.

Common Pitfalls

The biggest pitfall is rebuilding graphs or models repeatedly without clearing the associated framework state. That produces steady memory growth even if every session is closed.

Another pitfall is assuming GPU memory shown by monitoring tools must be a leak. TensorFlow often keeps reserved memory for reuse.

Developers also forget that notebooks and interactive shells are long-lived processes. Old variables stay alive unless you delete them or restart the kernel.

Finally, if memory grows across TensorFlow versions with the same code, test a smaller reproduction. Real leaks do happen, but many reports turn out to be graph accumulation or allocator behavior rather than a session close failure.

Summary

  • 'sess.close() releases session resources, but not every graph or Python reference around them.'
  • Rebuilding graphs inside loops is a common cause of memory growth.
  • Use tf.keras.backend.clear_session() and, when appropriate, reset_default_graph() to clear framework state.
  • GPU memory retention is often allocator caching, not a true leak.
  • Check the surrounding Python objects and long-lived process state before blaming the session alone.

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.