如何进一步简化CUDA_VISIBLE_DEVICES使用？求更优雅的别名替代方案

1. 动态GPU选择bash函数 #

直接在.bashrc中定义函数，既支持自动选空闲GPU，也支持手动指定编号，彻底替代硬编码别名：

# 用法：gpu [GPU数量 | GPU编号列表]
# 示例：gpu 2 → 自动选取2个利用率最低的空闲GPU
# 示例：gpu 0,3 → 固定使用GPU 0和3
gpu() {
    if [[ $# -eq 0 ]]; then
        echo "用法：gpu [GPU数量 | GPU编号列表]"
        return 1
    fi

    if [[ $1 =~ ^[0-9]+$ ]]; then
        # 自动筛选N个空闲GPU（基于nvidia-smi的利用率排序）
        local gpu_ids=$(nvidia-smi --query-gpu=index,utilization.gpu --format=csv,noheader,nounits | sort -n -k2 | head -n $1 | awk '{print $1}' | paste -sd "," -)
        export CUDA_VISIBLE_DEVICES=$gpu_ids
        echo "已设置CUDA_VISIBLE_DEVICES=$gpu_ids"
    else
        # 手动指定GPU编号
        export CUDA_VISIBLE_DEVICES=$1
        echo "已设置CUDA_VISIBLE_DEVICES=$1"
    fi
}

保存后执行source ~/.bashrc生效，后续直接用gpu 2或gpu 1,3即可快速设置。

2. 临时启动命令前缀 #

如果不想修改全局环境变量，可写一个前缀函数，直接在训练命令前指定GPU：

# 用法：g [GPU编号] 你的训练命令
# 示例：g 0,1 python train.py
g() {
    local gpus=${1:?请指定GPU编号，例如g 0,1}
    shift
    CUDA_VISIBLE_DEVICES=$gpus "$@"
}

每次启动训练时直接用g 0,1 python train.py，用完不影响全局环境变量。

3. 交互式可视化选择 #

先安装gpustat（pip install gpustat），再写一个交互式函数，看状态后再选GPU：

gpu_select() {
    gpustat
    read -p "请输入要使用的GPU编号（逗号分隔）：" gpus
    export CUDA_VISIBLE_DEVICES=$gpus
    echo "已设置CUDA_VISIBLE_DEVICES=$gpus"
}

执行gpu_select会先展示所有GPU的负载、显存使用情况，输入编号后自动完成设置，适合不确定空闲GPU的场景。

4. 可维护的预设组合 #

如果有固定常用的GPU组合，用数组存储比硬编码别名更易维护：

# 预设常用GPU组合
declare -A GPU_SETS=(
    ["single"]="4"
    ["double"]="2,5"
    ["triple"]="2,5,6"
    ["quad"]="0,4,6,7"
)

# 用法：gpu_set quad
gpu_set() {
    local set_name=$1
    if [[ -z ${GPU_SETS[$set_name]} ]]; then
        echo "可用组合：${!GPU_SETS[@]}"
        return 1
    fi
    export CUDA_VISIBLE_DEVICES=${GPU_SETS[$set_name]}
    echo "已设置CUDA_VISIBLE_DEVICES=${GPU_SETS[$set_name]}"
}

后续用gpu_set quad就能快速切换到预设的4卡组合，要修改组合只需更新GPU_SETS数组即可。

内容的提问来源于stack exchange，提问作者zheyuanWang