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()

参考文章