pynvml:获取和监控 NVIDIA GPU 设备的信息

发布时间 2023-09-28 14:12:01作者: 阿瑞娜

pynvml 是一个 Python 包,用于与 NVIDIA 的管理库 NVML(NVIDIA Management Library)进行交互,以获取和监控 NVIDIA GPU 设备的信息。它提供了一组函数和类,可以查询 GPU 设备的状态、显存使用情况、温度、功耗等信息。

要使用 pynvml 包,首先需要安装 NVIDIA 驱动程序和 CUDA 工具包。然后,可以使用以下步骤在 Python 中使用 pynvml

  1. 确保已安装 NVIDIA 驱动程序和 CUDA 工具包。你可以从 NVIDIA 官方网站下载并安装它们。
  2. 在 Python 环境中安装 pynvml 包。可以使用以下命令使用 pip 安装:
pip install pynvml
  1. 在 Python 脚本中导入 pynvml 模块:
import pynvml
  1. 初始化 pynvml
pynvml.nvmlInit()
  1. 使用 pynvml 提供的函数和类来获取 GPU 设备信息。例如,要获取 GPU 数量和每个 GPU 的名称,可以使用以下代码:
device_count = pynvml.nvmlDeviceGetCount()
for i in range(device_count):
    handle = pynvml.nvmlDeviceGetHandleByIndex(i)
    gpu_name = pynvml.nvmlDeviceGetName(handle)
    print(f"GPU {i}: {gpu_name}")
  1. 在使用完 pynvml 后,记得调用 pynvml.nvmlShutdown() 来清理和释放资源:
pynvml.nvmlShutdown()

这只是 pynvml 包的基本用法示例,你可以根据自己的需求进一步探索和利用其功能。

请注意,pynvml 只能在支持 NVIDIA GPU 的系统上正常工作,并且需要正确配置的驱动程序和 CUDA 工具包。