gpu轮询显存，自动跑程序脚本

import os
import time

gpu_device = 0
cmd = "CUDA_VISIBLE_DEVICES=0 python train.py"


def gpu_info(gpu_index):
    info = os.popen('nvidia-smi|grep %').read().split('\n')[gpu_index].split('|')
    memory = int(info[2].split('/')[0].strip()[:-3])
    return memory

# print(gpu_info(0))


while True:
    memory = gpu_info(gpu_device)
    if memory < 1000:
        break
    time.sleep(30)
    print("waiting | gpu ", str(gpu_device), " mem is ", memory)


os.system(cmd)

 

每30s轮询一次gpu_device 的显存，当显存低于1000M的时候，立刻运行cmd程序。