How to accurately track GPU VRAM usage?

54 views Asked by At

I want to know how much GPU VRAM each line of program will occupy. When I input 224x224x3 image to VGG16, the total parameters are 138,364,968. Each parameter can store dtype float32, then total VRAM usage that the code model.to("cuda") run should be spend 138364986*4(bytes)/1024**2=527.820534 MB.

When I use this code:

def get_gpu_memory_usage():
    result = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,noheader,nounits'])
    return int(result.decode().strip())

Catch each line like this:

model.to(device)
line1_memory = get_gpu_memory_usage() - initial_gpu_memory
line_memory_usage.append(line1_memory)

This line model.to(device) consumes 1125MB of VRAM usage. I think it should consume 527.820534 MB. Do I miss something?

0

There are 0 answers