Query PyTorch memory utilisation

torch_query_memory.py

#!/usr/bin/env python3

import os
from subprocess import check_output

def get_pmem():
    pmem = check_output(["sh", "-c", f"pmap -d {os.getpid()}|tail -n 1"], encoding="latin1").split()[3]
    return pmem

def get_gpumem():
    pid = os.getpid()
    csv=check_output(["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv"], encoding="latin1").split("\n")
    for pid, val in [x.split(",") for x in csv]:
        if pid == str(os.getpid()):
            return val
    return -1


if __name__ == "__main__":
    import torch
    x=torch.rand(3,3, device="cuda:0")
    print(torch.__version__, get_pmem(), get_gpumem())