不需要每次跑代码之前再检查空闲gpu了
import os
import numpy as np
def choose_gpu():
os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
memory_gpu=[int(x.split()[2]) for x in open('tmp','r').readlines()]
max_free = np.argmax(memory_gpu)
print("choose gpu %d free %d MiB"%(max_free, memory_gpu[max_free]))
os.environ['CUDA_VISIBLE_DEVICES']=str(max_free)
os.system('rm tmp')
if __name__ == "__main__":
choose_gpu()
参考自：https://blog.csdn.net/u011961856/article/details/77884946
网友评论