526互联

Python通过终端 cmd 获取GPU相关信息

发布时间 2023-04-28 13:21:35作者: ~逍遥子~
# 获取GPU相关信息
def get_gpu_info():
    try:
        cmd_out = subprocess.check_output('nvidia-smi --query-gpu=name,memory.used,memory.total --format=csv,noheader',
                                          shell=True)
        gpu_info = cmd_out.decode().strip().split('\n')
        gpu_info = [info.split(', ') for info in gpu_info]
        return gpu_info
    except subprocess.CalledProcessError as e:
        print("Error while invoking nvidia-smi: ", e)
        return None

# 打印 GPU 型号及占用情况
def print_gpu_usage():
    gpu_info = get_gpu_info()
    if gpu_info:
        total_memory = 0
        used_memory = 0
        for name, used, total in gpu_info:
            used_memory += int(used.strip().split()[0])
            total_memory += int(total.strip().split()[0])
            memory_usage_percent = round(used_memory / total_memory * 100, 2)
            print(f"GPU: {name.strip()}, Memory used: {used.strip()}, Memory total: {total.strip()}"
                  f", Memory usage: {memory_usage_percent}%")