python查看GPU容量

最新推荐文章于 2023-12-26 09:13:49 发布

师大云帆

最新推荐文章于 2023-12-26 09:13:49 发布

阅读量286

点赞数

分类专栏：深度学习文章标签： python 开发语言人工智能

本文链接：https://blog.csdn.net/qq_44996943/article/details/128258897

版权

深度学习专栏收录该内容

6 篇文章 0 订阅

订阅专栏

from pynvml import *
def check_gpu_mem_usedRate():
    max_rate = 0.0
    # while True:
    # pip install nvidia-ml-py
    nvidia_dict = {
        "state": True,
        "nvidia_version": "",
        "nvidia_count": 0,
        "gpus": []
    }
    try:
        nvmlInit()
        nvidia_dict["nvidia_version"] = nvmlSystemGetDriverVersion()
        nvidia_dict["nvidia_count"] = nvmlDeviceGetCount()
        for i in range(nvidia_dict["nvidia_count"]):
            handle = nvmlDeviceGetHandleByIndex(i)
            memory_info = nvmlDeviceGetMemoryInfo(handle)
            gpu = {
                "gpu_name": nvmlDeviceGetName(handle),
                "total": memory_info.total,
                "free": memory_info.free,
                "used": memory_info.used,
                "temperature": f"{nvmlDeviceGetTemperature(handle, 0)}℃",
                "powerStatus": nvmlDeviceGetPowerState(handle)
            }
            nvidia_dict['gpus'].append(gpu)
    except NVMLError as _:
        nvidia_dict["state"] = False
    except Exception as _:
        nvidia_dict["state"] = False
    finally:
        try:
            nvmlShutdown()
        except:
            pass
    info = nvidia_dict
    # print(info)
    used = info['gpus'][0]['used']  #B
    tot = info['gpus'][0]['total'] #B
    print('GPU0 used: %.1f MB'% (used/1024/1024), 'tot:%.1f MB'% (tot/1024/1024), '使用率：%.4f' % (used/tot))