武汉悠牛网网站建设加强服务保障 满足群众急需需求

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

pip install pynvml

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

import pynvmlpynvml.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')