GPU memory overhead for PyTorch

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)))