python实现自动查找并使用空闲GPU

查看GPU memory,并将结果保存在tmp中:

os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')

读取gpu memory:

memory_gpu=[int(x.split()[2]) for x in open('tmp','r').readlines()]

求剩余memory最多的显卡号,并设置CUDA_VISIBLE_DEVICES为该显卡:

os.environ['CUDA_VISIBLE_DEVICES']=str(np.argmax(memory_gpu))
os.system('rm tmp')

这里显卡1剩余memory最多,因此CUDA_VISIBLE_DEVICES=1,即使用1号显卡.

运行代码,终端输入nvidia-smi,即可查看显卡使用信息.