需要安装pynvml库。
下载地址为:https://pypi.org/project/nvidia-ml-py/#history
pip安装的命令为:
pip install nvidia-ml-py
代码如下:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2019/1/19 17:40
# @Author : xiaodai
import pynvml
pynvml.nvmlInit()
handle = len(pynvml.nvmlDeviceGetHandleByIndex(0))# 这里的0是GPU id
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total) #第二块显卡总的显存大小
print(meminfo.used)#这里是字节bytes,所以要想得到以兆M为单位就需要除以1024**2
print(meminfo.free) #第二块显卡剩余显存大小
print(pynvml.nvmlDeviceGetCount())#显示有几块GPU