Python使用pynvml查看GPU信息

前言

NVIDIA为自己家的显卡提供了NVML（英伟达显卡管理库）和构建在其上的nvidia-smi(显卡系统管理界面)，用户可以很方便的在系统命令窗口查询显卡信息和当前工作状况。

在Python中同样可以利用相应的package实现以上功能，这个好用的package就是pynvml库，可以通过它来获取NVIDIA显卡信息。

安装

直接通过pip安装：

pip install nvidia-ml-py

或者根据所使用的python版本安装对应包：

pip install nvidia-ml-py2 #python2

pip install nvidia-ml-py3 #python3

使用

import pynvml #导包

NUM_EXPAND = 1024 * 1024

"""GPU消耗情况"""
pynvml.nvmlInit() #初始化
#print("Driver: ", pynvml.nvmlSystemGetDriverVersion()) #显示驱动信息

all_gpu_used = []#用于记录所有GPU内存使用信息的list
need_record_pid = 19289 #需要记录的进程PID

gpuDeviceCount = pynvml.nvmlDeviceGetCount()#获取Nvidia GPU块数
for i in range(gpuDeviceCount):
    handle = pynvml.nvmlDeviceGetHandleByIndex(i)#获取GPU i的handle，后续通过handle来处理
    #info = pynvml.nvmlDeviceGetMemoryInfo(handle)#通过handle获取GPU i的信息
    ## gpu_memory_total = info.total #GPU i的总显存
    #gpu_memory_used = info.used / NUM_EXPAND #转为MB单位
    #all_gpu_used.append(gpu_memory_used) #添加进list
    
    ###还可以直接针对pid的gpu消耗进行统计
    info_list = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)#获取所有GPU上正在运行的进程信息
    info_list_len = len(info_list)
    gpu_memory_used = 0
    if info_list_len > 0:#0表示没有正在运行的进程
        for info_i in info_list:
            if info_i.pid == need_record_pid:#如果与需要记录的pid一致
                gpu_memory_used += info_i.usedGpuMemory / NUM_EXPAND #统计某pid使用的总显存
    all_gpu_used.append(gpu_memory_used)
pynvml.nvmlShutdown() #最后关闭管理工具

注：更多函数和功能可以参考NVIDIA官方GPU管理和部署文档。

参考资料

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