gpu轮询显存，自动跑程序

Posted 2021-09-13 立冬以东

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)