Introduction

Monitoring GPU performance during TensorFlow model training is essential for identifying bottlenecks, preventing out-of-memory errors, and ensuring your hardware is fully utilized. Without monitoring, you might train for hours with the GPU sitting at 10% utilization because the data pipeline is the bottleneck. TensorFlow provides built-in profiling tools, and NVIDIA offers command-line and programmatic interfaces to track GPU metrics in real time.

Using nvidia-smi for Real-Time Monitoring

The simplest approach is NVIDIA's nvidia-smi command-line tool, which reports GPU utilization, memory usage, temperature, and power consumption.

1# Snapshot of current GPU state
2nvidia-smi
3
4# Continuous monitoring every 1 second
5nvidia-smi --loop=1
6
7# Watch specific metrics with custom format
8nvidia-smi --query-gpu=timestamp,name,utilization.gpu,utilization.memory,memory.used,memory.total,temperature.gpu --format=csv --loop=1
9
10# Log to a CSV file for later analysis
11nvidia-smi --query-gpu=timestamp,utilization.gpu,memory.used --format=csv --loop=1 > gpu_log.csv

timestamp, name, utilization.gpu [%], utilization.memory [%], memory.used [MiB], memory.total [MiB], temperature.gpu
2025-01-15 10:30:01, NVIDIA A100, 95 %, 42 %, 32000 MiB, 40960 MiB, 62

A GPU utilization below 80% during training usually indicates a data pipeline bottleneck, not a GPU limitation.

TensorFlow Profiler (TensorBoard)

TensorFlow's built-in profiler provides detailed GPU traces showing kernel execution, memory allocation, and data transfer overhead.

1import tensorflow as tf
2
3# Create a TensorBoard callback with profiling enabled
4log_dir = "logs/profile"
5tensorboard_callback = tf.keras.callbacks.TensorBoard(
6    log_dir=log_dir,
7    profile_batch=(10, 20)  # Profile batches 10 through 20
8)
9
10model.fit(
11    train_dataset,
12    epochs=10,
13    callbacks=[tensorboard_callback]
14)

# Launch TensorBoard to view the profile
tensorboard --logdir=logs/profile
# Navigate to the "Profile" tab in the browser

The profiler shows a timeline of GPU operations, highlights idle periods, and provides recommendations for improving performance. The "Overview" page summarizes GPU utilization percentage and step time breakdown (input pipeline vs. computation).

Programmatic GPU Monitoring with GPUtil

For monitoring within your training script, the GPUtil library provides a Python interface to GPU metrics.

1import GPUtil
2import threading
3import time
4
5def monitor_gpu(interval=2):
6    """Print GPU stats at regular intervals during training."""
7    while getattr(monitor_gpu, "running", True):
8        gpus = GPUtil.getGPUs()
9        for gpu in gpus:
10            print(f"GPU {gpu.id}: {gpu.name} | "
11                  f"Load: {gpu.load * 100:.0f}% | "
12                  f"Memory: {gpu.memoryUsed:.0f}/{gpu.memoryTotal:.0f} MB | "
13                  f"Temp: {gpu.temperature}°C")
14        time.sleep(interval)
15
16# Start monitoring in a background thread
17monitor_gpu.running = True
18monitor_thread = threading.Thread(target=monitor_gpu, daemon=True)
19monitor_thread.start()
20
21# Run training
22model.fit(train_dataset, epochs=10)
23
24# Stop monitoring
25monitor_gpu.running = False

Custom Keras Callback for GPU Logging

1import tensorflow as tf
2import GPUtil
3
4class GPUMonitorCallback(tf.keras.callbacks.Callback):
5    def on_epoch_end(self, epoch, logs=None):
6        gpus = GPUtil.getGPUs()
7        for gpu in gpus:
8            print(f"Epoch {epoch} — GPU {gpu.id}: "
9                  f"Utilization {gpu.load * 100:.0f}%, "
10                  f"Memory {gpu.memoryUsed:.0f}/{gpu.memoryTotal:.0f} MB")
11
12    def on_train_begin(self, logs=None):
13        gpus = GPUtil.getGPUs()
14        for gpu in gpus:
15            print(f"Training on GPU {gpu.id}: {gpu.name}, "
16                  f"{gpu.memoryTotal:.0f} MB total memory")
17
18model.fit(train_dataset, epochs=10, callbacks=[GPUMonitorCallback()])

Checking TensorFlow GPU Configuration

Before training, verify that TensorFlow detects your GPU and is configured correctly.

1import tensorflow as tf
2
3# List available GPUs
4print("GPUs available:", tf.config.list_physical_devices('GPU'))
5
6# Check if TensorFlow is using GPU
7print("Built with CUDA:", tf.test.is_built_with_cuda())
8
9# Enable memory growth to avoid allocating all GPU memory at once
10gpus = tf.config.list_physical_devices('GPU')
11for gpu in gpus:
12    tf.config.experimental.set_memory_growth(gpu, True)

Common Pitfalls

• Low GPU utilization due to data pipeline bottleneck: If nvidia-smi shows low GPU utilization during training, the GPU is waiting for data. Use tf.data pipelines with .prefetch(tf.data.AUTOTUNE) and .map(..., num_parallel_calls=tf.data.AUTOTUNE) to overlap data loading with computation.
• Profiling too many batches: Setting profile_batch to a large range generates enormous trace files and slows training. Profile 5-20 batches to capture representative behavior without excessive overhead.
• TensorFlow allocating all GPU memory upfront: By default, TensorFlow reserves all GPU memory at process start. This prevents other processes from using the GPU and can cause OOM errors on multi-tenant systems. Enable set_memory_growth(True) to allocate memory incrementally.
• Ignoring memory fragmentation: Even when total GPU memory is sufficient, fragmentation can cause allocation failures. Monitor memory.used vs memory.total in nvidia-smi and watch for spikes during large batch operations.
• Not distinguishing GPU compute utilization from memory bandwidth utilization: nvidia-smi reports both utilization.gpu (SM occupancy) and utilization.memory (memory controller activity). High memory utilization with low compute utilization indicates memory-bound operations that might benefit from mixed-precision training (tf.keras.mixed_precision).

Summary

• Use nvidia-smi --loop=1 for quick real-time GPU monitoring from the terminal
• Use TensorFlow Profiler with TensorBoard for detailed operation-level traces and optimization recommendations
• Use GPUtil or custom Keras callbacks for programmatic monitoring within training scripts
• Enable set_memory_growth(True) to prevent TensorFlow from reserving all GPU memory
• Low GPU utilization usually indicates a data pipeline bottleneck — fix with tf.data prefetching and parallel mapping