使用场景:
我建立了一系列的tmux会话,他们的名字是gpu0,gpu1,gpu2等等。我调参的时候,每个程序占用一个gpu,我需要在程序里面写明白使用gpu 0,
os.environ['CUDA_VISIBLE_DEVICES']=f"{gpu}"
或者是
CUDA_VISIBLE_DEVICES=0 python train.py
然后把这个程序放到tmux名字是gpu0的会话里面用。这里存在一个冗余的地方,就是tmux已经指定了gpu id,程序再指定一次。
现在我通过下面的方案,去掉了修改程序内gpu ID的步骤。
通过os.system()执行下面的语句,能够得到当前tmux会话的名称,众所周知,这样得到的结果直接print了,无法用python获取。所以,我将结果保存到文本文件里面,然后用python读取temp.txt这个文件。
tmux display-message -p '#S'
tmux display-message -p '#S' >temp.txt
然后我就可以根据获取到的gpu id,动态的修改程序所使用的id。
os.environ['CUDA_VISIBLE_DEVICES']=f"{gpu}"
需要注意的几点:
- 在非tmux会话下运行,也会显示最后一次启动tmux会话的名称(如果有人有其他方法解决,欢迎留言)
- 修改gpu ID的语句要在使用gpu之前;



