Ubuntu/Linux下自动获取最佳GPU编号的脚本

最新推荐文章于 2026-07-25 16:18:12 发布

原创最新推荐文章于 2026-07-25 16:18:12 发布 · 1.4k 阅读

2 ·

本内容遵循CC 4.0 BY-SA版权协议

收录于

机器学习

Pytorch

PyTorch中如果直接用方法tensor.cuda()或tensor.device("cuda")将变量放到gpu上时默认将数据放到cuda:0上，而没有办法选择拥有最大空余存储的GPU上。因此写了一个脚本用来在Linux系统下自动获取拥有最大空余存储的GPU的对应ID（假设4张卡，默认ID为0到3）。

系统要求：Linux系统，且安装了nvidia-smi

首先给出实验设备的系统和环境：

# 查看操作系统
$ head -n 1 /etc/issue
Ubuntu 18.04.5 LTS \n \l

# 查看nvidia driver版本
$ cat /proc/driver/nvidia/version
NVRM version: NVIDIA UNIX x86_64 Kernel Module  465.27  Thu Apr 22 23:21:03 UTC 2021
GCC version:  gcc version 7.5.0 (Ubuntu 7.5.0-3ubuntu1~18.04)

文章目录

步骤

第一步

网上查到是可以直接用nvidia-smi加上--query-gpu参数进行查询。完整用法可以参考Useful nvidia-smi Queries
。

这里我选择的指令是

$ nvidia-smi --query-gpu=memory.free --format=csv

即以csv的格式输出memory.free的信息，输出顺序就是从cuda:0到cuda:3的顺序：

memory.free [MiB]
5126 MiB
7259 MiB
7259 MiB
7058 MiB

第二步

然后对这个返回进行一定处理，包括将所有换行符\n和空格' '删除，所用的指令是

tr -d '\n'

和

tr -d ' '

这步之后得到的结果为

$ nvidia-smi --query-gpu=memory.free --format=csv | tr -d '\n' | tr -d ' '
memory.free[MiB]5110MiB7259MiB7259MiB7058MiB

第三步

将prefix "memory.free\[MiB\]"删除，用到的指令是

$ sed -e "s/^$prefix//"

其中变量prefix="memory.free\[MiB\]"。

注意: 不能用prefix="memory.free[MiB]" !!!

此时得到的结果：

$ nvidia-smi --query-gpu=memory.free --format=csv | tr -d '\n' | tr -d ' ' | sed -e "s/^$prefix//"
5110MiB7259MiB7259MiB7058MiB

第四步

此时应该将所有的MiB替换为逗号','作为分割符，指令是

$ mem_free_list="${mem_free//MiB/,}"

其中mem_free就是第三步得到的结果5110MiB7259MiB7259MiB7058MiB。
由此得到

$ prefix="memory.free\[MiB\]"
# echo $prefix
$ mem_free=$(nvidia-smi --query-gpu=memory.free --format=csv | tr -d '\n' | tr -d ' ' | sed -e "s/^$prefix//")
$ mem_free_list="${mem_free//MiB/,}"
$ echo $mem_free_list
5110,7259,7259,7058,

第五步

从第四步中得到了大致pattern为num1,num2,num3,num4,的字符串，可以结合 Internal Field Separator（IFS）和for循环来找到free memory最大的CUDA编号。

oldIFS=$IFS
IFS=','
idx=0
max_idx=-1
max_mem=0
for cur_mem in $mem_free_list
do
# echo $idx $cur_mem

if (( max_mem  < cur_mem ))
    then
    max_mem=$cur_mem
    max_idx=$idx
fi 

let idx++
done

IFS=$oldIFS
echo "Max memory: $max_mem; CUDA ID: $max_idx"