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

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