Skip to content

Instantly share code, notes, and snippets.

@ovuruska

ovuruska/measure_gpu.py

Last active September 23, 2024 16:35
Show Gist options
  • Save ovuruska/55c6c50c9f60eb14c1136025f0b0c75e to your computer and use it in GitHub Desktop.
Save ovuruska/55c6c50c9f60eb14c1136025f0b0c75e to your computer and use it in GitHub Desktop.
Download ZIP
Measure Max GPU Usage: Python Decorator for CUDA Memory Tracking
Raw
measure_gpu.py
def measure_max_gpu_usage(func):
from torch.cuda import max_memory_allocated, reset_peak_memory_stats
def wrapper(*args, **kwargs):
reset_peak_memory_stats()
result = func(*args, **kwargs)
max_memory = max_memory_allocated() / 1024 / 1024 / 1024
print(f"Maximum GPU Usage: {max_memory:.3f} GB")
return result
return wrapper
############ Example usage ############
@measure_max_gpu_usage
def example_func():
pass
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment