单机多卡空闲GPU自动检测Pytorch多卡训练脚本
本文分享了一个Bash脚本，它能够自动检测空闲的GPU，并根据用户的需求分配相应数量的GPU以运行特定的命令或脚本。
脚本代码
#!/bin/bash
# 亮红色和亮绿色文本的ANSI转义序列
light_red='\e[91m'
light_green='\e[92m'
no_color='\e[0m'
# 检查是否有参数输入
if [ -z "$1" ]; then
echo -e "${light_red}Error: nproc_per_node is required. Please input a value.${no_color}"
exit 1
fi
nproc_per_node=$1
shift # 移除第一个参数，使得"$@"只包含除nproc_per_node之外的参数
# 检查nproc_per_node是否大于0
if ! [[ "$nproc_per_node" =~ ^[0-9]+$ ]] || [ "$nproc_per_node" -le 0 ]; then
echo -e "${light_red}Error: nproc_per_node must be a positive integer.${no_color}"
exit 1
fi
# 使用nvidia-smi查询空闲的GPU
free_gpus=($(nvidia-smi --query-gpu=index,gpu_name,utilization.gpu,utilization.memory,memory.free --format=csv,noheader,nounits | awk -F ', ' '{if($3 == 0 && $4 == 0) print $1}'))
# 检查是否有足够的空闲GPU
if [ ${#free_gpus[@]} -lt $nproc_per_node ]; then
echo -e "${light_red}Error: Not enough free GPUs. Found ${#free_gpus[@]}, require $nproc_per_node.${no_color}"
exit 1
fi
# 将所有空闲的GPU编号转换为以逗号分隔的字符串
all_free_gpus=$(IFS=,; echo "${free_gpus[*]}")
# 选择前nproc_per_node个空闲的GPU
selected_gpus=$(IFS=,; echo "${free_gpus[*]:0:$nproc_per_node}")
# 输出所有空闲的GPU编号
echo -e "${light_green}All free GPUs: ${all_free_gpus}${no_color}"
# 输出选择的GPU编号
echo -e "${light_green}Selected GPUs: ${selected_gpus}${no_color}"
# 构建并以绿色文本运行命令
cmd="CUDA_VISIBLE_DEVICES=$selected_gpus python -m torch.distributed.launch --nproc_per_node=$nproc_per_node --use_env $@"
echo -e "${light_green}Running command: $cmd${no_color}"
eval $cmd
功能简述
该脚本的核心功能是自动检测系统中空闲的GPU，并根据用户定义的需要的卡数（通过nproc_per_node参数）分配GPU资源。如果系统中没有足够的空闲GPU，脚本将输出错误信息并终止执行。
使用步骤
-
创建脚本文件：首先，您需要创建一个新的脚本文件。可以使用任何文本编辑器来创建文件，例如使用
vim。创建文件的命令如下：vim gpu_allocator.sh在打开的编辑器中，复制上文中的脚本代码，并保存退出。
-
赋予执行权限：为了让脚本文件可执行，您需要修改其权限。这可以通过下面的命令来完成：
chmod +x gpu_allocator.sh这个命令会赋予脚本文件执行权限。
-
执行脚本：现在，脚本已经准备好可以执行了。使用下面的格式运行您的脚本：
./gpu_allocator.sh <nproc_per_node> <其他命令><nproc_per_node>：这是一个必需的参数，指定您希望使用的GPU数量（整数）。<其他命令>：这是脚本将要执行的其他命令或脚本，例如运行Python脚本或其他需要GPU的应用。
示例
假设您想要使用2个GPU来运行一个Python训练
脚本train.py，您可以这样执行上述脚本：
./gpu_allocator.sh 2 python train.py --batch-size 64
在这个示例中，脚本会自动寻找2个空闲的GPU，如果找到，就会使用这两个GPU来执行train.py脚本，其中--batch-size 64是传递给train.py的一个参数。
原文地址：https://blog.csdn.net/qq_36396406/article/details/137656585
免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！