本文最后更新于:2022年7月4日 上午
深度学习过程中经常需要和GPU打交道,需要在训练之前确定GPU工作状态。本文介绍python中的pynvml包的使用方法,用来获取上述信息。
安装
直接通过pip安装:
1
| pip install nvidia-ml-py
|
或者根据所使用的python版本安装对应包:
1 2
| pip install nvidia-ml-py2 pip install nvidia-ml-py3
|
使用方法
1 2 3 4 5 6
| import pynvml
NUM_EXPAND = 1024 * 1024
pynvml.nvmlInit()
|
1 2 3
| pynvml.nvmlSystemGetDriverVersion() -> b'441.22'
|
1 2 3
| pynvml.nvmlDeviceGetCount() -> 1
|
1 2
| handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
1 2 3
| pynvml.nvmlDeviceGetName(handle) -> b'GeForce GTX 1660'
|
1 2
| info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
1 2 3 4 5
| info.total -> 6442450944 info.total / NUM_EXPAND -> 6144.0
|
1 2 3 4 5
| info.used -> 383299584 info.used / NUM_EXPAND -> 365.54296875
|
踩坑
- 报错 : pynvml.NVMLError_LibraryNotFound: NVML Shared Library Not Found
坑已填:解决 pynvml 报错 NVML Shared Library Not Found / WinError 126
参考资料
文章链接:
https://www.zywvvd.com/notes/study/deep-learning/pynvml-gpu-info-get/pynvml-gpu-info-get/