Does calling tf.keras.backend.clear_session() fully release memory when using subclassed Models with custom training loops?

clear_session() memory release with subclassed Models and custom train_step

Calling tf.keras.backend.clear_session() resets the default TensorFlow graph and removes Keras‑namespace references to model layers, variables, and the optimizer object. For a subclassed Model that implements a custom train_step, the method does not automatically garbage‑collect the Python model instance, its trainable weights, or optimizer variable slots that remain attached to the instance.

Confirmed facts

• clear_session() removes the default graph, so any tf.Variable or optimizer state that was only referenced through the Keras namespace is no longer reachable from that namespace.
• Optimizer variable slots (e.g., momentum, second‑moment estimates) and any custom attributes defined in your train_step stay attached to the model object unless you explicitly delete or reassign them.
• In eager execution, memory reclamation depends on Python’s garbage collector; the session reset alone does not guarantee immediate heap freeing.
• Under tf.function, the compiled graph holds persistent references to the model’s variables; destroying the session removes the graph, but the model object must go out of scope or be deleted to free memory.

Steps to achieve full reclamation in this case

1. After your custom training loop, run tf.keras.backend.clear_session().
2. Immediately afterwards, delete the model instance: del model.
3. If you use callbacks (e.g., ModelCheckpoint, TensorBoard), remove or reset them before the call, or ensure they do not hold references to the model’s weights or optimizer state.
4. Explicitly reset optimizer state if your loop re‑initializes or re‑uses the same optimizer across iterations: optimizer.reset_states() or create a fresh optimizer instance each iteration.
5. Run import gc; gc.collect() to force Python garbage collection, then verify heap size with tracemalloc or tf.profiler.

Execution‑mode difference

With eager code, clear_session() followed by gc.collect() typically suffices if no hidden references persist. Under tf.function, you must also ensure the function‑wrapped model and its variables are out of scope; otherwise the compiled graph retains them until the process restarts.

Do any of your callbacks or external objects retain a reference to the subclassed model or its optimizer after the loop ends? Answering this will determine whether the above steps are sufficient or if additional callback‑reset logic is required.