How to Check GPU Memory Usage in PyTorch Easily

Use torch.cuda.memory_allocated() to check the current GPU memory used by tensors and torch.cuda.memory_reserved() to see the total reserved GPU memory by PyTorch. These functions help monitor GPU memory usage during model training or inference.

Syntax

PyTorch provides simple functions to check GPU memory usage:

• torch.cuda.memory_allocated(device=None): Returns the current GPU memory used by tensors in bytes.
• torch.cuda.memory_reserved(device=None): Returns the total GPU memory reserved by the caching allocator in bytes.
• device parameter is optional and can specify which GPU device to check (e.g., 0 for the first GPU).

python

import torch

# Check memory allocated on default GPU
allocated = torch.cuda.memory_allocated()

# Check memory reserved on default GPU
reserved = torch.cuda.memory_reserved()

print(f"Allocated memory: {allocated} bytes")
print(f"Reserved memory: {reserved} bytes")

Example

This example shows how to check GPU memory usage before and after creating a tensor on the GPU.

python

import torch

def print_gpu_memory():
    allocated = torch.cuda.memory_allocated()
    reserved = torch.cuda.memory_reserved()
    print(f"Allocated memory: {allocated} bytes")
    print(f"Reserved memory: {reserved} bytes")

if torch.cuda.is_available():
    print("Before tensor allocation:")
    print_gpu_memory()

    # Create a tensor on GPU
    x = torch.randn(1000, 1000, device='cuda')

    print("After tensor allocation:")
    print_gpu_memory()
else:
    print("CUDA is not available on this system.")

Output

Before tensor allocation:
Allocated memory: 0 bytes
Reserved memory: 0 bytes
After tensor allocation:
Allocated memory: 4000000 bytes
Reserved memory: 4194304 bytes

Common Pitfalls

Common mistakes when checking GPU memory usage in PyTorch:

• Not checking if CUDA is available before calling memory functions, which causes errors on CPU-only machines.
• Confusing memory_allocated() (memory used by tensors) with memory_reserved() (memory reserved by PyTorch allocator).
• Not resetting or clearing GPU cache, which can make memory usage appear higher than actual tensor usage.

Always use torch.cuda.empty_cache() if you want to clear unused cached memory.

python

import torch

if torch.cuda.is_available():
    print("Memory before empty_cache:")
    print(torch.cuda.memory_reserved())
    torch.cuda.empty_cache()
    print("Memory after empty_cache:")
    print(torch.cuda.memory_reserved())

Output

Memory before empty_cache:
4194304
Memory after empty_cache:
0

Quick Reference

Summary of key PyTorch GPU memory functions:

Function	Description
torch.cuda.memory_allocated(device=None)	Returns current GPU memory used by tensors in bytes.
torch.cuda.memory_reserved(device=None)	Returns total GPU memory reserved by PyTorch caching allocator in bytes.
torch.cuda.empty_cache()	Releases unused cached memory back to the GPU.
torch.cuda.max_memory_allocated(device=None)	Returns the maximum GPU memory allocated by tensors since the program started.
torch.cuda.reset_max_memory_allocated(device=None)	Resets the peak memory statistics.

Key Takeaways

Use torch.cuda.memory_allocated() to check actual GPU memory used by tensors.

Use torch.cuda.memory_reserved() to see total memory reserved by PyTorch's allocator.

Always check torch.cuda.is_available() before querying GPU memory to avoid errors.

Clear unused GPU memory with torch.cuda.empty_cache() to get accurate memory readings.

Monitor peak memory usage with torch.cuda.max_memory_allocated() for optimization.