查看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,即可查看显卡使用信息.