Profile GPU Usage in Python with NVML

Here is the article you requested, written for PythonSkillset.com.

Profiling Python GPU Usage with NVML

You're running a machine learning model, and it feels slow. You know the GPU is involved, but you have no idea if it's actually working hard or just sitting idle. Guessing won't fix performance. You need numbers.

That's where NVML comes in. The NVIDIA Management Library gives you direct access to your GPU's vital signs. And with the pynvml Python wrapper, you can hook into those stats from your own scripts.

Let's walk through a practical way to profile your Python code and see exactly what the GPU is doing.

Why Just "top" Won't Cut It

nvidia-smi is a great command-line tool. But it gives you a snapshot of the entire system. When you're running a Python script, you need to know what that specific process is doing. You need to sample utilization while the code runs, not after.

Using NVML from Python lets you: - Track GPU utilization per process. - Log memory usage over time. - Pinpoint bottlenecks (compute-bound vs. memory-bound).

This turns guessing into diagnostics.

Getting Started with pynvml

First, install the library. It's a simple pip install.

pip install nvidia-ml-py3

Then, you initialize the library in your Python code. This connects to the NVIDIA driver.

import pynvml

pynvml.nvmlInit()

You almost always want to know how many GPUs you have.

device_count = pynvml.nvmlDeviceGetCount()
for i in range(device_count):
    handle = pynvml.nvmlDeviceGetHandleByIndex(i)
    name = pynvml.nvmlDeviceGetName(handle)
    print(f"GPU {i}: {name}")

This gives you a baseline. You know which GPU to watch.

Sampling Utilization During a Function

The core of profiling is taking measurements around your code. Let's build a simple context manager that measures GPU utilization for a block of code.

import time
from contextlib import contextmanager

@contextmanager
def profile_gpu_utilization(device_index=0):
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)

    # Get starting info
    start_util = pynvml.nvmlDeviceGetUtilizationRates(handle)
    start_mem = pynvml.nvmlDeviceGetMemoryInfo(handle)

    start_time = time.time()
    yield
    end_time = time.time()

    # Get ending info
    end_util = pynvml.nvmlDeviceGetUtilizationRates(handle)
    end_mem = pynvml.nvmlDeviceGetMemoryInfo(handle)

    elapsed = end_time - start_time

    print(f"Elapsed time: {elapsed:.2f}s")
    print(f"GPU Utilization: {end_util.gpu}% (start: {start_util.gpu}%)")
    print(f"Memory Utilization: {end_util.memory}%")
    print(f"Memory Used: {end_mem.used / 1024**2:.0f} MB / {end_mem.total / 1024**2:.0f} MB")

    # Risk of stale reading since we only capture start/end
    # For better results, sample in a separate thread

You use it like this:

with profile_gpu_utilization(0):
    # Your GPU-intensive code here
    for epoch in range(5):
        # training loop
        pass

The Problem with Snapshot Sampling

The code above only checks utilization at two points. If your GPU code runs for a few seconds, you might miss the peak. A better approach is to sample continuously in a background thread.

Here is a more robust version that logs utilization every 0.5 seconds.

import threading
import time
import queue

class GPUProfiler:
    def __init__(self, device_index=0, interval=0.5):
        self.device_index = device_index
        self.interval = interval
        self._queue = queue.Queue()
        self._running = False
        self._thread = None
        pynvml.nvmlInit()
        self._handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)

    def _sample_loop(self):
        while self._running:
            util = pynvml.nvmlDeviceGetUtilizationRates(self._handle)
            mem = pynvml.nvmlDeviceGetMemoryInfo(self._handle)
            self._queue.put({
                'timestamp': time.time(),
                'gpu_util': util.gpu,
                'mem_used': mem.used
            })
            time.sleep(self.interval)

    def start(self):
        self._running = True
        self._thread = threading.Thread(target=self._sample_loop, daemon=True)
        self._thread.start()

    def stop(self):
        self._running = False
        if self._thread:
            self._thread.join()

    def get_samples(self):
        samples = []
        while not self._queue.empty():
            samples.append(self._queue.get())
        return samples

You use it like this inside your training script.

profiler = GPUProfiler(0)
profiler.start()

# Your training loop
for epoch in range(10):
    # train model
    time.sleep(1.1)  # Simulate work

profiler.stop()
samples = profiler.get_samples()

# Analyze samples
avg_util = sum(s['gpu_util'] for s in samples) / len(samples) if samples else 0
print(f"Average GPU Utilization: {avg_util:.1f}%")

What to Look For in the Data

Once you have a log of utilization over time, you can spot patterns.

• Low utilization (< 50%): Your GPU is waiting on data. This usually means your data loading pipeline is the bottleneck. Look into faster I/O, data prefetching, or larger batch sizes.
• High utilization (> 90%): Your GPU is compute-bound. You are pushing it to its limits. This is usually a good sign, but if it's sustained at 100%, you may be hitting thermal limits or throttling.
• Memory spikes: If memory usage grows over time, you may have a memory leak. If it hits the limit, your process will crash.

A Real World Example

At PythonSkillset, we had a user who was training a transformer model. The code looked straightforward. But the training was twice as slow as expected. They ran the profiler from above.

The samples showed GPU utilization bouncing between 30% and 45%. That was a red flag. They checked the data loader. The bottleneck was a slow image decoding step. They used Pillow in a single thread. Switching to torchdata with multiple workers fixed it. Utilization jumped to 85%, and training time dropped by 40%.

Without the profiler, they would have been guessing. With it, they had proof.

Final Tips

• Always close the NVML connection at the end of your script: pynvml.nvmlShutdown().
• Be careful with the sampling interval. Too fast (under 0.1s) can add overhead. Too slow (over 1s) can miss spikes.
• For long-running training jobs, log samples to a file and plot them afterward. matplotlib is your friend here.

Profiling isn't optional. If you don't measure, you're flying blind. NVML gives you the instruments. Now go see what your GPU is actually doing.