TensorFlow
GPU Memory
Programmatically Determine
Machine Learning
Neural Networks

how to programmatically determine available GPU memory with tensorflow?

Master System Design with Codemia

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

Introduction

Programmatically checking "available GPU memory" with TensorFlow is trickier than it sounds because there are two different questions you might mean. One is how much memory TensorFlow's allocator is currently using. The other is how much free memory the GPU has overall at the device level. TensorFlow can report allocator statistics directly, but device-wide free memory is usually easier to query through NVIDIA's management library.

TensorFlow Allocator Memory Info

In TensorFlow 2, the most direct built-in API is tf.config.experimental.get_memory_info.

python
1import tensorflow as tf
2
3print(tf.config.list_physical_devices("GPU"))
4info = tf.config.experimental.get_memory_info("GPU:0")
5print(info)

The returned dictionary typically includes fields such as:

  • 'current for memory currently used by TensorFlow's allocator'
  • 'peak for peak allocator usage'

This is useful when you want to understand TensorFlow's own memory behavior during model creation or training.

It is not the same as asking the GPU driver how much memory the whole device still has free for every process on the system.

Device-Wide Free Memory With NVML

For actual free GPU memory on NVIDIA devices, NVML is usually the more accurate tool. In Python, the common wrapper is pynvml.

python
1import pynvml
2
3pynvml.nvmlInit()
4handle = pynvml.nvmlDeviceGetHandleByIndex(0)
5mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
6
7print("total:", mem.total)
8print("used:", mem.used)
9print("free:", mem.free)

This reports device-level totals from the NVIDIA driver, which is closer to what most people mean by available GPU memory.

If your TensorFlow program shares the GPU with other processes, NVML gives a better picture of global pressure than TensorFlow's allocator stats alone.

Using Both Together

The most informative setup is often to read both views.

python
1import tensorflow as tf
2import pynvml
3
4print(tf.config.experimental.get_memory_info("GPU:0"))
5
6pynvml.nvmlInit()
7handle = pynvml.nvmlDeviceGetHandleByIndex(0)
8mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
9print({"total": mem.total, "used": mem.used, "free": mem.free})

This lets you compare:

  • memory managed by TensorFlow
  • memory occupied on the device overall

That distinction matters because TensorFlow may reserve memory or share the GPU with other libraries and processes.

Controlling Memory Growth

TensorFlow can reserve GPU memory aggressively unless configured otherwise. Enabling memory growth often makes diagnostics easier.

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4if gpus:
5    tf.config.experimental.set_memory_growth(gpus[0], True)

With memory growth enabled, TensorFlow allocates GPU memory more incrementally instead of grabbing most of it upfront.

That does not directly answer the memory query question, but it changes how the reported numbers behave and often prevents confusion when monitoring allocations.

A Practical Helper Function

Here is a small helper that returns both TensorFlow allocator stats and NVML stats when available.

python
1import tensorflow as tf
2
3
4def gpu_memory_snapshot(device="GPU:0"):
5    result = {"tensorflow": tf.config.experimental.get_memory_info(device)}
6
7    try:
8        import pynvml
9        pynvml.nvmlInit()
10        index = int(device.split(":")[1])
11        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
12        mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
13        result["nvml"] = {"total": mem.total, "used": mem.used, "free": mem.free}
14    except Exception as exc:
15        result["nvml_error"] = str(exc)
16
17    return result
18
19
20print(gpu_memory_snapshot())

This is useful in training scripts, experiments, and debugging notebooks.

Common Pitfalls

A common mistake is assuming TensorFlow's get_memory_info reports total free GPU memory for the whole system. It usually reports allocator information, not a full device-wide view.

Another mistake is forgetting that TensorFlow may reserve memory early unless memory growth is enabled. That can make the GPU appear more occupied than expected.

Developers also sometimes query GPU state before TensorFlow has initialized the device, then compare it to values gathered later under a different memory policy.

Finally, remember that these techniques are most practical for NVIDIA GPUs. The NVML approach is vendor-specific.

Summary

  • TensorFlow can report its own allocator memory with tf.config.experimental.get_memory_info.
  • That built-in API is not the same as global device-wide free memory.
  • For NVIDIA GPUs, NVML via pynvml is the usual way to get total, used, and free device memory.
  • Using both TensorFlow stats and NVML gives the clearest picture.
  • Memory growth changes TensorFlow allocation behavior and can make diagnostics easier.
  • Be explicit about whether you want allocator usage or whole-device availability.

Course illustration
Course illustration

All Rights Reserved.