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')