TensorFlow
GPU memory management
deep learning
Python
machine learning

How to prevent tensorflow from allocating the totality of a GPU memory?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

TensorFlow often reserves most of a GPU by default once it initializes that device. If you want it to allocate memory only as needed, the usual supported fix is to enable memory growth before any GPU work starts. If you want a hard cap instead, configure a logical device with a memory limit.

Why TensorFlow Grabs So Much GPU Memory

TensorFlow prefers predictable GPU allocation behavior. Reserving a large block early can reduce fragmentation and avoid repeated allocation overhead during training. That is useful for single-purpose training jobs, but it is inconvenient when:

  • several processes need the same GPU
  • you want interactive experimentation
  • another framework also needs VRAM
  • the machine runs notebooks, inference, and training together

So the problem is not unusual. It is just TensorFlow's default behavior.

Use Memory Growth

The most common fix is set_memory_growth:

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices('GPU')
4if gpus:
5    for gpu in gpus:
6        tf.config.experimental.set_memory_growth(gpu, True)

With memory growth enabled, TensorFlow starts small and expands its allocations when more GPU memory is actually needed.

This must happen before TensorFlow initializes the GPU. If your code has already created tensors or models on the GPU, changing the setting is too late.

Example With Error Handling

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices('GPU')
4
5try:
6    if gpus:
7        for gpu in gpus:
8            tf.config.experimental.set_memory_growth(gpu, True)
9        print("Memory growth enabled")
10    else:
11        print("No GPU detected")
12except RuntimeError as err:
13    print(f"Configuration failed: {err}")

The RuntimeError typically means the GPU runtime was already initialized before you tried to configure it.

Use a Hard Memory Limit When Needed

Sometimes gradual growth is not enough. If you want TensorFlow to stay below a specific VRAM limit, create a logical device:

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices('GPU')
4if gpus:
5    tf.config.set_logical_device_configuration(
6        gpus[0],
7        [tf.config.LogicalDeviceConfiguration(memory_limit=4096)]
8    )

This example limits the first GPU to about 4 GB for TensorFlow's use. That can be helpful in shared servers or when you want reproducible resource ceilings.

Pick the Right Strategy

Use memory growth when:

  • one process can use variable memory over time
  • you want TensorFlow to coexist more politely with other workloads
  • you do not know the exact memory ceiling in advance

Use logical device limits when:

  • you want a strict cap
  • several jobs must divide the GPU intentionally
  • exceeding a quota would disrupt another workload

These settings solve different operational problems.

Do This Before Import Side Effects

Some libraries trigger GPU initialization as a side effect. To stay safe, put your GPU configuration near the start of the program, before model creation and ideally before helper modules build any TensorFlow objects.

A good structure is:

python
1import tensorflow as tf
2
3def configure_gpu():
4    gpus = tf.config.list_physical_devices('GPU')
5    for gpu in gpus:
6        tf.config.experimental.set_memory_growth(gpu, True)
7
8configure_gpu()
9
10# Build models only after configuration
11model = tf.keras.Sequential([
12    tf.keras.layers.Dense(32, activation='relu'),
13    tf.keras.layers.Dense(1)
14])

Environment and Tooling Still Matter

Even correct TensorFlow code will not behave well if the CUDA stack, GPU driver, or runtime libraries are mismatched. If memory settings appear ignored, verify that:

  • TensorFlow sees the GPU
  • the driver and CUDA stack match your TensorFlow build expectations
  • another process is not already occupying the memory

Those issues are separate from TensorFlow's allocation policy.

Common Pitfalls

The most common mistake is calling set_memory_growth after TensorFlow has already initialized the GPU. Another is assuming memory growth imposes a hard ceiling. It does not. It only changes allocation behavior over time.

Developers also sometimes expect the setting to solve every out-of-memory problem. If the model genuinely needs more VRAM than the device has, growth will not save it. In that case you need a smaller batch size, a smaller model, mixed precision where appropriate, or a larger GPU.

Summary

  • TensorFlow often reserves most GPU memory by default.
  • Enable memory growth to make it allocate VRAM gradually.
  • Use logical device configuration if you need a strict memory cap.
  • Apply these settings before TensorFlow initializes the GPU.
  • If memory is still exhausted, the model may simply need more optimization or a larger device.

Course illustration
Course illustration

All Rights Reserved.