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

为为为什么

发布于 2022-08-05 03:06:42

1.6K00

代码可运行

文章被收录于专栏：又见苍岚又见苍岚

运行总次数：0

代码可运行

深度学习过程中经常需要和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 报错 NVML Shared Library Not Found / WinError 126

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2020年11月6日，如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自作者个人站点/博客前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体同步曝光计划，欢迎热爱写作的你一起参与！

登录后参与评论

0 条评论

热度