抢占GPU显卡跑程序--实用脚本

本文主要是介绍抢占GPU显卡跑程序--实用脚本，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

只需修改两行，将脚本复制到程序位置，预定义要抢占的显卡号，写入自己要启动的程序指令

#!/bin/bash
#=====================只需关注==================
# 要检查的 GPU 列表，使用逗号分隔，如 "4,5,6,7"
GPUS_TO_CHECK="4,5,6,7"
# 要启动的程序
PROGRAM="nohup sh train.sh"
#=====================只需关注==================# 显存剩余百分比阈值
MEMORY_THRESHOLD=80# 检查显卡显存空闲的函数
check_gpu_memory() {local gpu_id=$1# 获取指定 GPU 的总显存和已用显存local memory_total=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits -i $gpu_id)local memory_used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i $gpu_id)# 计算剩余显存百分比local memory_free_percentage=$((100 * (memory_total - memory_used) / memory_total))# 如果显存剩余百分比大于等于阈值，则认为显卡空闲if [ "$memory_free_percentage" -ge "$MEMORY_THRESHOLD" ]; thenreturn 0elsereturn 1fi
}# 计数空闲 GPU 的数量
free_gpu_count=0# 遍历 GPU 列表并检查显存状态
for gpu_id in $(echo $GPUS_TO_CHECK | tr ',' ' '); doif check_gpu_memory $gpu_id; thenfree_gpu_count=$((free_gpu_count + 1))fi
done# 检查是否有至少两张 GPU 空闲
if [ "$free_gpu_count" -ge 2 ]; thenecho "At least two GPUs are free. Starting the program..."CUDA_VISIBLE_DEVICES=$(echo $GPUS_TO_CHECK | cut -d',' -f1,2) $PROGRAM &exit 0
elseecho "Less than two GPUs are free. Exiting..."exit 1
fi