aa.sh文件(小白文章,有错误麻烦指出,谢谢)
#!/bin/bash
var=1 #控制最外层while中断
while [ ${var} -eq 1 ]
do
count=0 #定义一个变量,=两旁不能有空格,以字符串形式识别
for i in $(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)
#nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits:查询GPU的内存使用信息,以CSV格式返回(无表头和单位)
do
if [ ${i} -lt 3000 ] #如何GPU卡i的内存使用小于3000,两个[]有空格
then
echo 'GPU'$count' is avaiable' #向终端输出GPUi是有空闲
source /home/ubuntu/miniconda3/etc/profile.d/conda.sh #启动conda脚本
conda activate pytorch #启动conda脚本后才能使用conda命令
/home/ubuntu/ #切换程序目录,如果此处不切换,所执行程序中目录为当前.sh文件目录
python train17.py --gpus 'cuda:'$count #执行train17.py,并向里传送GPU使用显卡参数
var=0
break
fi
count=$(($count+1)) #当前显卡数+1
done
done
trian17.py文件代码
def train(arguments):
device = arguments.gpus
print(device)
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='CNN Seg Training Function')
parser.add_argument("--gpus", type=str, default = 'cuda:0')#接受使用显卡参数
args = parser.parse_args()
train(args)
$bash aa.sh #在linux终端执行.sh文件