Python使用pynvml查看GPU信息

参考：

https://blog.csdn.net/TracelessLe/article/details/107405544

====================================================

发现了一个python版本的可以查看和管理NVIDIA显卡的工具pynvml，安装方法：

pip install nvidia-ml-py

在这个库的基础上自己写了一个小程序，用来监控GPU的使用情况：

import psutil
import pynvml #导包


UNIT = 1024 * 1024


pynvml.nvmlInit() #初始化
gpuDeriveInfo = pynvml.nvmlSystemGetDriverVersion()
print("Drive版本: ", str(gpuDeriveInfo, encoding='utf-8')) #显示驱动信息


gpuDeviceCount = pynvml.nvmlDeviceGetCount()#获取Nvidia GPU块数
print("GPU个数：", gpuDeviceCount )


for i in range(gpuDeviceCount):
    handle = pynvml.nvmlDeviceGetHandleByIndex(i)#获取GPU i的handle，后续通过handle来处理

    memoryInfo = pynvml.nvmlDeviceGetMemoryInfo(handle)#通过handle获取GPU i的信息

    gpuName = str(pynvml.nvmlDeviceGetName(handle), encoding='utf-8')

    gpuTemperature = pynvml.nvmlDeviceGetTemperature(handle, 0)

    gpuFanSpeed = pynvml.nvmlDeviceGetFanSpeed(handle)

    gpuPowerState = pynvml.nvmlDeviceGetPowerState(handle)

    gpuUtilRate = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
    gpuMemoryRate = pynvml.nvmlDeviceGetUtilizationRates(handle).memory

    print("第 %d 张卡："%i, "-"*30)
    print("显卡名：", gpuName)
    print("内存总容量：", memoryInfo.total/UNIT, "MB")
    print("使用容量：", memoryInfo.used/UNIT, "MB")
    print("剩余容量：", memoryInfo.free/UNIT, "MB")
    print("显存空闲率：", memoryInfo.free/memoryInfo.total)
    print("温度：", gpuTemperature, "摄氏度")
    print("风扇速率：", gpuFanSpeed)
    print("供电水平：", gpuPowerState)
    print("gpu计算核心满速使用率：", gpuUtilRate)
    print("gpu内存读写满速使用率：", gpuMemoryRate)
    print("内存占用率：", memoryInfo.used/memoryInfo.total)

    """
    # 设置显卡工作模式
    # 设置完显卡驱动模式后，需要重启才能生效
    # 0 为 WDDM模式，1为TCC 模式
    gpuMode = 0     # WDDM
    gpuMode = 1     # TCC
    pynvml.nvmlDeviceSetDriverModel(handle, gpuMode)
    # 很多显卡不支持设置模式，会报错
    # pynvml.nvml.NVMLError_NotSupported: Not Supported
    """

    # 对pid的gpu消耗进行统计
    pidAllInfo = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)#获取所有GPU上正在运行的进程信息
    for pidInfo in pidAllInfo:
        pidUser = psutil.Process(pidInfo.pid).username()
        print("进程pid：", pidInfo.pid, "用户名：", pidUser, 
            "显存占有：", pidInfo.usedGpuMemory/UNIT, "Mb") # 统计某pid使用的显存


pynvml.nvmlShutdown() #最后关闭管理工具

输出：

Drive版本:  470.63.01
GPU个数： 4
第 0 张卡： ------------------------------
显卡名： NVIDIA GeForce RTX 3090
内存总容量： 24268.3125 MB
使用容量： 14168.625 MB
剩余容量： 10099.6875 MB
显存空闲率： 0.41616768780276747
温度： 59 摄氏度
风扇速率： 51
供电水平： 2
gpu计算核心满速使用率： 46
gpu内存读写满速使用率： 46
内存占用率： 0.5838323121972325
进程pid： 68028 用户名： xxxxxx  显存占有： 14165.0 Mb
第 1 张卡： ------------------------------
显卡名： NVIDIA GeForce RTX 3090
内存总容量： 24268.3125 MB
使用容量： 3.625 MB
剩余容量： 24264.6875 MB
显存空闲率： 0.9998506282626779
温度： 32 摄氏度
风扇速率： 30
供电水平： 8
gpu计算核心满速使用率： 0
gpu内存读写满速使用率： 0
内存占用率： 0.0001493717373220738
第 2 张卡： ------------------------------
显卡名： NVIDIA GeForce RTX 3090
内存总容量： 24268.3125 MB
使用容量： 3.625 MB
剩余容量： 24264.6875 MB
显存空闲率： 0.9998506282626779
温度： 33 摄氏度
风扇速率： 30
供电水平： 8
gpu计算核心满速使用率： 0
gpu内存读写满速使用率： 0
内存占用率： 0.0001493717373220738
第 3 张卡： ------------------------------
显卡名： NVIDIA GeForce RTX 3090
内存总容量： 24268.3125 MB
使用容量： 3.625 MB
剩余容量： 24264.6875 MB
显存空闲率： 0.9998506282626779
温度： 33 摄氏度
风扇速率： 30
供电水平： 8
gpu计算核心满速使用率： 0
gpu内存读写满速使用率： 0
内存占用率： 0.0001493717373220738

注意的是上面代码需要安装第三方依赖库：psutil

相关资料：

[1] NVIDIA官方GPU管理和部署文档
[2] 【Python管理GPU】pynvml工具的安装与使用
[3] pypi下的nvidia-ml-py包

https://pythonhosted.org/nvidia-ml-py/

https://github.com/gpuopenanalytics/pynvml

==================================================