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

本文最后更新于:2021年11月1日 中午

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

安装

直接通过pip安装:

1
pip install nvidia-ml-py

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

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

使用方法

1
2
3
4
5
6
# 导入包
import pynvml
# 每MB包含的字节数
NUM_EXPAND = 1024 * 1024
# 初始化
pynvml.nvmlInit()
1
2
3
#GPU驱动信息
pynvml.nvmlSystemGetDriverVersion()
-> b'441.22'
1
2
3
#获取Nvidia GPU块数
pynvml.nvmlDeviceGetCount()
-> 1
1
2
#获取GPU 0 的handle
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
1
2
3
# 设备名称
pynvml.nvmlDeviceGetName(handle)
-> b'GeForce GTX 1660'
1
2
#通过handle获取GPU 0 的信息
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
1
2
3
4
5
#GPU 0 的总显存
info.total
-> 6442450944
info.total / NUM_EXPAND
-> 6144.0
1
2
3
4
5
#GPU 0 使用的显存
info.used
-> 383299584
info.used / NUM_EXPAND
-> 365.54296875
1
2
#关闭 pynvml
pynvml.nvmlShutdown()

踩坑

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

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

参考资料