1、安装

pip方式安装：

pip install nvidia-ml-py

根据python版本制定2/3：

pip install nvidia-ml-py2  # python2
pip install nvidia-ml-py3  # python3

源码安装：

#下载链接：http://pypi.python/pypi/nvidia-ml-py/
sudo python setup.py install

2、使用

2.1、初始化、析构

import pynvml
pynvml.nvmlInit() # 初始化
... # 函数调用
pynvml.nvmlShutdown() # 最后要关闭管理工具

2.2、获取驱动版本号

pynvml.nvmlSystemGetDriverVersion()
=>
b'426.00' # 版本号426

2.3、获取显卡个数

pynvml.nvmlDeviceGetCount()
=>
1 # 1块显卡

2.4、获取显卡句柄

根据GPU id获取显卡句柄,id为0,1,2,3...

handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)

2.5、获取显卡名字-型号

pynvml.nvmlDeviceGetName(handle)
=>
b'GeForce GTX 1060 6GB'

2.6、获取显卡内存信息Info

meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total)  # 显卡总的显存大小,6442450944Bit
print(meminfo.used)  # 显存使用大小,4401950720Bit
print(meminfo.free)  # 显卡剩余显存大小,2040500224Bit
meminfo.used / 1024 / 1024 = 4198 M

2.7、获取显卡温度、风扇、电源

print("Temperature is %d C"%nvmlDeviceGetTemperature(handle,0))
print("Fan speed is "nvmlDeviceGetFanSpeed(handle))
print("Power ststus",nvmlDeviceGetPowerState(handle))
=>
Temperature is 34 C
Fan speed is  0
Power ststus 8

3、使用实例

#!/usr/bin/env python
# coding: utf-8

import pynvml

def query_by_id(gpu_id):
    '''
       查询gpuid下显卡使用显存情况
       :param gpu_id: 有效的gpuid，数字类型，无效gpuid则报错.
       :return: 使用显存大小，单位M
    '''

    pynvml.nvmlInit()
    print(pynvml.nvmlDeviceGetCount())  # 显示有几块GPU
    handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)

    print(meminfo.total)  # 显卡总的显存大小
    print(meminfo.used)  # 显存使用大小
    print(meminfo.free)  # 显卡剩余显存大小
    return meminfo.used / 1024 / 1024


def main():
    gpu_id = 0
    memory_used = query_by_id(gpu_id)
    print("gpu %d used memory is %d M." % (gpu_id, memory_used))
    pass

4、参考资源

nvmlDeviceXXX有一系列函数可以调用，包括了NVML的大多数函数。具体可以参考：

https://docs.nvidia/deploy/nvml-api/group__nvmlDeviceQueries.html#group__nvmlDeviceQueries

nvidia_smi:https://github/ultrabug/py3status/blob/master/py3status/modules/nvidia_smi.py

https://pythonhosted/nvidia-ml-py/

传送门：python应用资源总目录