pynvml 是一个 Python 包,用于与 NVIDIA 的管理库 NVML(NVIDIA Management Library)进行交互,以获取和监控 NVIDIA GPU 设备的信息。它提供了一组函数和类,可以查询 GPU 设备的状态、显存使用情况、温度、功耗等信息。
文档
安装
需要安装 NVIDIA 驱动程序和 CUDA 工具包
pip install pynvml
示例
import pynvml
# 初始化 pynvml
pynvml.nvmlInit()
# 获取 GPU 数量
print(pynvml.nvmlDeviceGetCount())
# 获取GPU信息
for i in range(deviceCount):
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
gpu_name = pynvml.nvmlDeviceGetName(handle)
# 清理和释放资源
pynvml.nvmlShutdown()