Introduction

When people ask how to log GPU VRAM utilization in TensorFlow, they usually mean one of two different things. Sometimes they want TensorFlow's own view of memory usage during a training run, and sometimes they want the GPU's device-level memory and utilization numbers. Those are related, but they are not the same measurement, so the right solution depends on what you are trying to observe.

Use TensorFlow's Built-In Memory Stats First

TensorFlow exposes device memory statistics through tf.config.experimental.get_memory_info. This is the easiest way to inspect what TensorFlow is actively using on a GPU.

1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4if not gpus:
5    print("No GPU found")
6else:
7    stats = tf.config.experimental.get_memory_info("GPU:0")
8    print("Current bytes:", stats["current"])
9    print("Peak bytes:", stats["peak"])

This is useful for answering questions like:

• how much memory is my TensorFlow job using right now
• what was the peak memory usage during this run
• did a specific training step increase memory pressure

That makes it a good fit for model debugging and batch-size tuning.

Understand What TensorFlow Is Reporting

The important detail is that TensorFlow reports the memory it is actually using, not necessarily the total memory it has reserved from the GPU driver. That distinction matters because TensorFlow may allocate aggressively unless you change the default behavior.

If you want more predictable allocation behavior, enable memory growth before TensorFlow initializes the GPU:

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 grows usage as needed. That often makes logging easier to interpret, especially during local experimentation.

Measure Peaks Around a Specific Block of Code

Peak memory numbers become more useful when you reset them before the part of the program you care about.

1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4if gpus:
5    tf.config.experimental.reset_memory_stats("GPU:0")
6
7    x = tf.random.normal((4096, 4096))
8    y = tf.matmul(x, x)
9
10    stats = tf.config.experimental.get_memory_info("GPU:0")
11    print("Current bytes:", stats["current"])
12    print("Peak bytes:", stats["peak"])

This pattern is useful when you want to compare:

• one batch size versus another
• one model block versus another
• eager execution versus tf.function

It gives you a tight measurement window instead of one large peak from the whole program.

Log Memory During Training

You can also log TensorFlow memory usage from a callback:

1import tensorflow as tf
2import numpy as np
3
4
5class GpuMemoryLogger(tf.keras.callbacks.Callback):
6    def on_epoch_end(self, epoch, logs=None):
7        gpus = tf.config.list_physical_devices("GPU")
8        if gpus:
9            stats = tf.config.experimental.get_memory_info("GPU:0")
10            print(
11                f"epoch={epoch + 1} current={stats['current']} peak={stats['peak']}"
12            )
13
14
15x = np.random.random((128, 10)).astype("float32")
16y = np.random.random((128, 1)).astype("float32")
17
18model = tf.keras.Sequential(
19    [
20        tf.keras.layers.Input(shape=(10,)),
21        tf.keras.layers.Dense(32, activation="relu"),
22        tf.keras.layers.Dense(1),
23    ]
24)
25
26model.compile(optimizer="adam", loss="mse")
27model.fit(x, y, epochs=2, callbacks=[GpuMemoryLogger()], verbose=0)

That gives you TensorFlow-aware logging without leaving Python.

Use nvidia-smi for Device-Level VRAM and GPU Utilization

TensorFlow does not provide every GPU metric you may want. If you need device-level VRAM usage, total memory, or GPU utilization percentages, use nvidia-smi.

1import subprocess
2
3result = subprocess.run(
4    [
5        "nvidia-smi",
6        "--query-gpu=memory.used,memory.total,utilization.gpu",
7        "--format=csv,noheader,nounits",
8    ],
9    capture_output=True,
10    text=True,
11    check=True,
12)
13
14for index, line in enumerate(result.stdout.strip().splitlines()):
15    used, total, util = [part.strip() for part in line.split(",")]
16    print(f"gpu={index} used_mib={used} total_mib={total} util_percent={util}")

This is the better choice when you care about the GPU as a whole rather than only TensorFlow's allocator state. It is also helpful when multiple processes share the same device.

Which Approach Should You Use

Use TensorFlow's API when you want per-run memory insight from inside the training job. Use nvidia-smi when you want hardware-level monitoring or when you need to compare TensorFlow with other processes using the same GPU.

In practice, many teams use both:

• TensorFlow memory stats for debugging model behavior
• 'nvidia-smi for operational monitoring and dashboards'

That combination gives a clearer picture than either tool alone.

Common Pitfalls

One common mistake is assuming TensorFlow's memory stats are the same as total VRAM reserved by the process. They are not. TensorFlow reports current and peak memory it is actually using, which can differ from what the driver shows as allocated.

Another mistake is trying to enable memory growth after the GPU has already been initialized. TensorFlow requires that configuration to happen before tensors are created or the runtime touches the device.

Developers also sometimes ask TensorFlow for GPU utilization percentage. TensorFlow's memory APIs do not give you the full nvidia-smi view of device utilization. If you need that metric, query the GPU directly.

Finally, if your script must also work on CPU-only machines, guard GPU calls with tf.config.list_physical_devices("GPU") so the program fails gracefully instead of throwing a device error.

Summary

• Use tf.config.experimental.get_memory_info("GPU:0") for TensorFlow's current and peak GPU memory usage.
• Use reset_memory_stats() to measure the peak for a specific section of code.
• Enable memory growth early if you want TensorFlow to allocate GPU memory incrementally.
• Use nvidia-smi when you need device-level VRAM usage or GPU utilization percentages.
• TensorFlow memory stats and hardware-level VRAM metrics answer different questions, so choose accordingly.