Skip to content

Instantly share code, notes, and snippets.

@ebarsoum

ebarsoum/gpu_memory_overhead_pytorch.py

Created October 26, 2019 00:20
Show Gist options
  • Save ebarsoum/898cc72871e8defd84928b940532c11a to your computer and use it in GitHub Desktop.
Save ebarsoum/898cc72871e8defd84928b940532c11a to your computer and use it in GitHub Desktop.
Download ZIP
GPU memory overhead for PyTorch
Raw
gpu_memory_overhead_pytorch.py
import torch
import numpy as np
from pynvml.smi import nvidia_smi
nvsmi = nvidia_smi.getInstance()
def getGPUMemoryUsage(gpu_index=0):
return nvsmi.DeviceQuery("memory.used")["gpu"][gpu_index]['fb_memory_usage']['used']
gpu_index = 0
device = torch.device('cuda:{}'.format(gpu_index))
print("Before: used GPU Memory: {} MB".format(getGPUMemoryUsage(gpu_index)))
torch_tensor = torch.tensor(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32)).to(device)
print("Max memory cached: {} MB".format(torch.cuda.max_memory_cached() / (1024*1024)))
print("After: used GPU Memory: {} MB".format(getGPUMemoryUsage(gpu_index)))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment