✍14 pynvml模块的使用(GPU)

查看GPU信息

nvidia-smi

简单使用

# 导入模块
import  pynvml

# 每MB包含的字节数(用于后面单位换算)
NUM_EXPAND = 1024 * 1024

# 初始化
pynvml.nvmlInit()
'''
pynvml.nvml.NVMLError_LibraryNotFound: NVML Shared Library Not Found
1.出现这个问题可能是你的电脑没有显卡, pynvml源码里面会去寻找电脑显卡驱动安装目录, 没找到就报错
2.目录错误--> 解决方法--->
https://www.zywvvd.com/2020/11/06/deep_learning/pynvml_model_not_found/pynvml_model_not_found/
https://www.cnblogs.com/ShineLeBlog/p/11219781.html
'''

# 获取GPU驱动版本信息
version = pynvml.nvmlSystemGetDriverVersion()
print(version)

#获取Nvidia GPU个数
count = pynvml.nvmlDeviceGetCount()
print(count)

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

# 通过 handle 设备名称
dev_name = pynvml.nvmlDeviceGetName(handle)
print(dev_name)

# 通过 handle 获取uuid
uuid = pynvml.nvmlDeviceGetUUID(handle)
print(uuid)

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

# 获取GPU 0 的总显存
print(info.total)
print(f"{info.total/NUM_EXPAND}MB")

# 获取GPU 0 已经使用的显存
print(info.used)
print(f"{info.used/NUM_EXPAND}MB")

#关闭 pynvml
pynvml.nvmlShutdown()

测试

• 155的服务器上有显卡, 使用155的机器做测试

更多方法

请参考 : NVIDIA官方GPU管理和部署文档