安装：

#https://pypi.org/project/nvidia-ml-py/
pip install nvidia-ml-py
pip install nvidia-ml-py2
pip install nvidia-ml-py3

使用示例：

import pynvml

def log_gpu_memory():
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    info.free = round(info.free / 1024**2)
    info.used = round(info.used / 1024**2)
    print('[INFO]GPU memory free: {}, memory used: {}'.format(info.free, info.used))
    return info.used


def get_gpu_memory_total():
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    info.total = round(info.total / 1024**2)
    return info.total

#
pynvml.nvmlInit()
print("[INFO]Driver: ", pynvml.nvmlSystemGetDriverVersion()) 
setup_default_logging()
log_gpu_memory()
total_gpu_mem = get_gpu_memory_total()
#
deviceCount = pynvml.nvmlDeviceGetCount()
for i in range(deviceCount):
    handle = pynvml.nvmlDeviceGetHandleByIndex(i)
    print(f"[INFO]GPU {i}:", pynvml.nvmlDeviceGetName(handle))
    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    print("[INFO]Memory Total: ",info.total)
    print("[INFO]Memory Free: ",info.free)
    print("[INFO]Memory Used: ",info.used)
    print("[INFO]Temperature is %d C", pynvml.nvmlDeviceGetTemperature(handle,0))
    print("[INFO]Fan speed is ", pynvml.nvmlDeviceGetFanSpeed(handle))
    print("[INFO]Power ststus", pynvml.nvmlDeviceGetPowerState(handle))

#
pynvml.nvmlShutdown()