import subprocess
support_info = [
"memory.total",
"memory.free",
"memory.used",
"name", # 后三项还没
"utilization.gpu",
"utilization.memory",
]
[docs]
def get_gpu_info(s):
if s in support_info:
cmd = "nvidia-smi --query-gpu={} --format=csv,noheader".format(s)
ret = subprocess.getoutput(cmd)
return ret.split("\n")
else:
raise Exception("Not support")
[docs]
def get_gpu_total_memory(i=0):
mem = get_gpu_info("memory.total")
ret = mem[i].split(" ")[0]
return float(ret)
[docs]
def get_gpu_free_memory(i=0):
mem = get_gpu_info("memory.free")
ret = mem[i].split(" ")[0]
return float(ret)
[docs]
def get_gpu_used_memory(i=0):
mem = get_gpu_info("memory.used")
ret = mem[i].split(" ")[0]
return float(ret)