Python使用pynvml查看GPU信息

发现了一个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​​

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