python-GPU狀態-pynvml

python-pynvml

目錄

• python-pynvml
  • 簡介
  • 安裝
  • 使用
  • 獲取驅動版本號
  • 獲取顯卡句柄
    • 獲得GPU數量
    • 獲得GPU名稱
    • 獲取GPU 0 內存信息
    • 獲得GPU 進程信息
  • 案例
  • 其他方法

簡介

深度學習過程中經常需要和GPU打交道，需要在訓練之前確定GPU工作狀態。

Nvidia的顯卡提供了 NVML（英偉達顯卡管理庫）以及構建在其上的 nvidia-smi 顯卡系統管理界面)，可以方便的查詢顯卡的信息和工作狀況。在python中同樣可以利用pynvml庫來實現顯卡信息的獲取

官方指導文件

https://docs.nvidia.com/deploy/nvml-api/index.html

安裝

pip install nvidia-ml-py

pip install nvidia-ml-py3  # python3

使用

import pynvml

pynvml.nvmlInit() # 初始化
... # 函數調用
pynvml.nvmlShutdown() # 最后要關閉管理工具

獲取驅動版本號

#GPU驅動信息

pynvml.nvmlSystemGetDriverVersion()
=>
b'426.00' # 版本號426

獲取顯卡句柄

獲得GPU數量

#獲取Nvidia GPU塊數
pynvml.nvmlDeviceGetCount()
-> 1

獲得GPU名稱

handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)

name=pynvml.nvmlDeviceGetName(handle)

print(name)

->  b'GeForce GTX 1650 6GB'

獲取GPU 0 內存信息

# 每MB包含的字節數
NUM_EXPAND = 1024 * 1024

gpu_id=0
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)

gpu_Total = info.total   # 總顯存
gpu_Free = info.free 
gpu_Used = info.used 

print(gpu_Total)  # 顯卡總的顯存大小,6442450944Bit
print(gpu_Free)  # 顯存使用大小,4401950720Bit
print(gpu_Used)  # 顯卡剩余顯存大小,2040500224Bit

print(gpu_Total / NUM_EXPAND)
print(gpu_Free / NUM_EXPAND)
print(gpu_Used / NUM_EXPAND)

# meminfo.used / 1024 / 1024  
#    4198 M

獲得GPU 進程信息

獲得運行下的進程，進程線程大小

# 
gpu_id=0

handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
info_list = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)

for each_pidinfo in info_list:
	
    print(type(each_pidinfo))
    print(each_pidinfo.__dict__)   # 打印pid的屬性
    print(each_pidinfo.pid)

    

-> <class 'pynvml.nvml.nvmlFriendlyObject'>    
    
-> {'pid': 31566, 'usedGpuMemory': 1259339776, 'gpuInstanceId': 4294967295, 'computeInstanceId': 4294967295}
    
-> 31566
    

案例

#簡單使用
from pynvml import *
nvmlInit()     #初始化
print("Driver: "nvmlSystemGetDriverVersion())  #顯示驅動信息
#>>> Driver: 384.xxx

#查看設備
deviceCount = nvmlDeviceGetCount()
for i in range(deviceCount):
    handle = nvmlDeviceGetHandleByIndex(i)
    print("GPU", i, ":", nvmlDeviceGetName(handle))
#>>>
#GPU 0 : b'GeForce GTX 1080 Ti'
#GPU 1 : b'GeForce GTX 1080 Ti'

#查看顯存、溫度、風扇、電源
handle = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(handle)
print("Memory Total: ",info.total)
print("Memory Free: ",info.free)
print("Memory Used: ",info.used)

print("Temperature is %d C"%nvmlDeviceGetTemperature(handle,0))
print("Fan speed is "nvmlDeviceGetFanSpeed(handle))
print("Power ststus",nvmlDeviceGetPowerState(handle))


#最后要關閉管理工具
nvmlShutdown()


#nvmlDeviceXXX有一系列函數可以調用，包括了NVML的大多數函數。
#具體可以參考：https://docs.nvidia.com/deploy/nvml-api/group__nvmlDeviceQueries.html#group__nvmlDeviceQueries

其他方法

手動解析 nvidia-smi 命令

以下是GPU p40的命令解析

import subprocess
def get_p40_pid():
    shell_data = subprocess.check_output('nvidia-smi', shell=True)
    shell_data = str(shell_data)
    shell_data = \
        shell_data.split('|=============================================================================|')[-1]
    shell_data = \
        shell_data.split('+-----------------------------------------------------------------------------+')[0]
    shell_data = shell_data.split(r'\n')[1:-1]
    
    Pid_dicts = {}
    for text in shell_data:
        pidtext = text.split()
        
        pid_id = pidtext[2]
        pid_name = pidtext[4]
        pid_gpu_memory = pidtext[5][:-3]  # [:-3]去掉單位


        Pid_dicts[str(pid_id)] = {"pid_id": pid_id,
                                  "pid_name": pid_name,
                                  "pid_gpu_memory": pid_gpu_memory}
    return Pid_dicts