python监控显卡显存

在工作组中，经常会遇到多人争抢显卡的情况，僧多肉少，所以可以做个监控显卡的头部程序，然后再调取主程序训练模型。
我们借用pynvml，没有安装的，可以使用下面命令直接安装：

pip install pynvml

下面使用pynvml可以查看显卡总共多少显存，剩余多少显存，以及已用多少显存。

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 这里的0表示第一张显卡
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total/1024/1024/1024) # 输出第一张显卡的显存
print(meminfo.free/1024/1024/1024) # 输出第一张显卡的可用显存
print(meminfo.used/1024/1024/1024) # 输出第一张显卡的已用显存

我们如果想对显卡进行监控，可以使用下面的命令：

import pynvml

# 定义一个读取gpu显卡的函数，gpu为第几张显卡，free表示需要多少剩余显存
def test_gpu(gpu, free):
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(gpu)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)

    return meminfo.free/1024/1024/1024 < free

然后用while循环对显卡进行监控：

while test_gpu(1, 40):
	print("GPU memory < %dGB, wait 1 minut"%40)
    time.sleep(60)

当显卡的剩余显存大于40GB的时候，才可以执行下面的程序，否则将等待1分钟（可以根据自己情况设置）。
如果想要执行python程序文件，可以使用os库运行：

import os
os.system('python run.py')