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 工具包。