python 使用 pynvml 工具获取GPU工作状态

深度学习过程中经常需要和GPU打交道，需要在训练之前确定GPU工作状态。本文介绍python中的pynvml包的使用方法，用来获取上述信息。

安装

直接通过pip安装：

pip install nvidia-ml-py

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

pip install nvidia-ml-py2  # python2
pip install nvidia-ml-py3  # python3

使用方法

# 导入包
import pynvml 
# 每MB包含的字节数
NUM_EXPAND = 1024 * 1024
# 初始化
pynvml.nvmlInit()

#GPU驱动信息
pynvml.nvmlSystemGetDriverVersion()
-> b'441.22'

#获取Nvidia GPU块数
pynvml.nvmlDeviceGetCount()
-> 1

#获取GPU 0 的handle
handle = pynvml.nvmlDeviceGetHandleByIndex(0)

# 设备名称
pynvml.nvmlDeviceGetName(handle)
-> b'GeForce GTX 1660'

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

#GPU 0 的总显存
info.total 
-> 6442450944
info.total / NUM_EXPAND
-> 6144.0

#GPU 0 使用的显存
info.used
-> 383299584
info.used / NUM_EXPAND
-> 365.54296875

#关闭 pynvml
pynvml.nvmlShutdown() 

踩坑

• 报错 ： pynvml.NVMLError_LibraryNotFound: NVML Shared Library Not Found

坑已填：解决 pynvml 报错 NVML Shared Library Not Found / WinError 126

参考资料

• http://www.ngui.cc/51cto/show-33630.html