1.安装pynvml
pip install nvidia-ml-py
2.输入下面代码
import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)# 这里的0是GPU id
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total/1024/1024)
print(meminfo.used/1024/1024)
print(meminfo.free/1024/1024)
print(pynvml.nvmlDeviceGetCount())#显示有几块GPU