nvidia-ml-py：获取显卡的相关信息

有时候我们想要查询 GPU 的相关信息，这个时候要怎么做呢？首先有一款软件叫 GPU-Z，可以很方便的查询，不过我们这里介绍是通过 Python 的方式。另外这里先解释一个概念，显卡是由 GPU 加显存组成的，GPU 和显存的关系类似于 CPU 和内存的关系。

首先安装一个模块：pip install nvidia-ml-py，然后就可以使用了。

import pynvml

# 初始化
pynvml.nvmlInit()

# 显示 GPU 块数，我这里是 1 块
print(pynvml.nvmlDeviceGetCount())  # 1

# 获取 Driver 版本
print(pynvml.nvmlSystemGetDriverVersion())  # 456.7

# 获取 handler，索引 0 表示获取第一块
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
# 基于 handler 得到设备信息
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
# 显存的总大小
print(info.total / 1024 ** 2)  # 2048.0
# 已使用的大小
print(info.used / 1024 ** 2)  # 1536.5078125
# 剩余大小，total - used
print(info.free)  # 536338432

# 显卡的名字
print(pynvml.nvmlDeviceGetName(handle))  # GeForce GTX 1050

# 显卡温度，当前是 54 度
print(pynvml.nvmlDeviceGetTemperature(handle, 0))  # 54

# 显卡风扇的转速
print(pynvml.nvmlDeviceGetFanSpeed(handle))  # 49

# 供电水平
print(pynvml.nvmlDeviceGetPowerState(handle))  # 0

# GPU 计算核心满速使用率
print(pynvml.nvmlDeviceGetUtilizationRates(handle).gpu)  # 52

# GPU 内存读写满速使用率
print(pynvml.nvmlDeviceGetUtilizationRates(handle).memory)  # 40

# 获取显卡的工作模式，0 为 WDDM 模式、1 为 TCC 模式
# WDDM 模式：GPU 负责计算以及图像显示，打游戏显然是 WDDM 模式
# TCC 模式：GPU 只负责计算，不负责图像显示，机器学习的时候一般是 TCC 模式
print(pynvml.nvmlDeviceGetDriverModel(handle))  # [0, 0]
"""
# 可以设置显卡工作模式，如果报如下错误
# pynvml.NVMLError_NoPermission: Insufficient Permissions
# 那么以管理员的身份打开命令行，然后执行
pynvml.nvmlDeviceSetDriverModel(handle, 1)
# 当然，如果显卡不是 Tesla 和 Quadro，那么不支持设置工作模式，此时会抛出如下错误：
# pynvml.NVMLError_NotSupported: Not Supported
"""

# 获取运行在 GPU 上的所有进程的 pid 信息
pid_info_list = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
for pid_info in pid_info_list:
    """
    pid_info.pid: 此进程的pid，再基于 psutil.Process 可以得到更详细的信息
    pid_info.usedGpuMemory: 此进程使用的显存大小 
    """
    
# 关闭
pynvml.nvmlShutdown()

当然还可以获取更多详细信息，可以自己尝试一下。