How to clear GPU memory in tensorflow 2?

First, separate two things people both call "clearing GPU memory":

1. Preventing TensorFlow from reserving the whole GPU at startup → configure allocation (Options 1 & 2 below).
2. Releasing memory a model is already holding → reset the session (the section after).

Configuring allocation

In TensorFlow 2, you configure GPU memory with the tf.config.experimental.set_memory_growth method (enable memory growth), or tf.config.experimental.set_virtual_device_configuration (a hard limit).

Let's go through both options with detailed explanations and examples:

Option 1: Enable Memory Growth By enabling memory growth, TensorFlow will allocate memory on the GPU as needed and release it when no longer in use. This allows the GPU memory to be freed up automatically. Here's how you can do it:

import tensorflow as tf
 
# Enable memory growth
physical_devices = tf.config.list_physical_devices('GPU')
if physical_devices:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

Explanation:

1. First, import the tensorflow module.
2. Use tf.config.list_physical_devices('GPU') to get a list of available GPUs.
3. Check if any GPUs are available using if physical_devices: to avoid errors.
4. Set memory growth to True using tf.config.experimental.set_memory_growth for the first GPU in the list (assuming you have multiple GPUs, you can modify the index if needed).

By setting memory growth to True, TensorFlow will allocate GPU memory on an as-needed basis. This allows the memory to grow dynamically based on the requirements of your model.

Option 2: Limit GPU Memory Usage If you want to set a specific limit on GPU memory usage, you can use tf.config.experimental.set_virtual_device_configuration. Here's an example:

import tensorflow as tf
 
# Limit GPU memory usage
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # Set a limit of 2GB for GPU memory
        tf.config.experimental.set_virtual_device_configuration(
            gpus[0],
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])
    except RuntimeError as e:
        print(e)

Explanation:

1. Import the tensorflow module.
2. Get the list of available GPUs using tf.config.list_physical_devices('GPU').
3. Check if any GPUs are available using if gpus:.
4. Set the memory limit for the first GPU using tf.config.experimental.set_virtual_device_configuration.
   • In the example, we limit the GPU memory to 2GB by passing memory_limit=2048 to VirtualDeviceConfiguration.
   • You can modify the memory limit value as per your requirements.
   • If you have multiple GPUs, you can set memory limits for each of them by specifying the appropriate index and memory limit.

By setting a memory limit, TensorFlow will allocate only the specified amount of GPU memory. If the model's memory requirements exceed the limit, an error will be raised.

Remember to place these code snippets before creating any TensorFlow operations or models to ensure the GPU memory configuration takes effect.

These methods provide you with the flexibility to manage GPU memory according to your needs, either by allowing memory growth or setting memory limits. Choose the option that suits your requirements best.

Releasing memory held within a process

The options above are set-once configuration. But the most common real problem is different: you build many models in a loop (hyperparameter search, cross-validation) and each one leaves memory behind until you hit an OOM. To reset TensorFlow's state between iterations:

import gc
import tensorflow as tf
 
for params in search_space:
    model = build_model(params)
    model.fit(...)
 
    # release this iteration's graph/session state
    del model
    tf.keras.backend.clear_session()   # clears the Keras graph and session
    gc.collect()                       # drop lingering Python references

clear_session() clears the current Keras graph and frees the objects tied to it, which is what lets a long loop keep running. However, TensorFlow's GPU allocator holds onto the underlying device memory for the life of the process — nvidia-smi will still show it reserved. That's expected, not a leak.

The only guaranteed full reset

If you truly need the GPU memory returned to the OS, end the process. In a notebook that means restarting the kernel; in a script, run each heavy job as a separate subprocess so exiting reclaims everything:

# run one training job per subprocess so its GPU memory is fully reclaimed on exit
import subprocess
subprocess.run(["python", "train_one.py", "--config", "a.json"])

Which approach to use

Related guides

• AI/ML interview & engineering hub

Sources

• TensorFlow — Use a GPU (memory growth)
• tf.keras.backend.clear_session — TensorFlow API