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]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_memif (( max_mem  < cur_mem ))thenmax_mem=$cur_memmax_idx=$idx
fi let idx++
doneIFS=$oldIFS
echo "Max memory: $max_mem; CUDA ID: $max_idx"

执行结果为输出

Max memory: 7259; CUDA ID: 1

即可以在pytorch的启动脚本里加上指定CUDA编号为1。

总结

完整脚本链接： max_mem_gpu.sh，脚本中将执行过程写成了函数，方便调用。

对shell脚本编程不是很熟悉，因此用的是比较笨的方法。如果有更简洁的指令希望有大佬分享

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

linux 自动安装 get,linux下自动获取并安装软件包 apt-get 的命令介绍
命令作用 apt-cache search package 搜索包 apt-cache show package 获取包的 ...
linux备份mysql怎样操作,Linux下自动备份MySQL数据库详细操作步骤(转载)
环境说明操作系统:CentOS IP:192.168.150.214 Oracle数据库版本:Oracle11gR2 用户:root 密码:123456 端口:3306 数据库:ts_0.ts_1. ...
【转载】【重要】Ubuntu Linux 下 Ffmpeg 及 Mencoder 安装使用小结
近段时间在学习 Ffmpeg.Mencoder 音视频转换,大概在 Ubuntu 10.04.Ubuntu 10.10.Ubuntu Natty 测试版上安装和使用过.重点是针对低分辨率的普通手机视频 ...
Ubuntu Linux下与MOTO E2手机进行USB NET连接
买了台新手机MOTO E2,linux操作系统.几天来几乎都在为了能让E2能在UBUNTU下进行SyncML而忙乎.E2刷到了49P(屠龙刀),手机支持USB NET方式连接,在WINDOWS下装好驱 ...
Ubuntu Linux 下文件名乱码（无效的编码）的快速解决办法
Ubuntu Linux 下文件名乱码(无效的编码)的快速解决办法文件是在WIndows 下创建的,Windows 的文件名中文编码默认为GBK,而Linux中默认文件名编码为UTF8,由于编码不 ...
Linux下自动备份Oracle数据库并删除指定天数前的备份
说明: Oracle数据库服务器操作系统:CentOS IP:192.168.0.198 端口:1521 SID:orcl Oracle数据库版本:Oracle11gR2 具体操作: 1.root用 ...
在UBUNTU LINUX下搭建 ANDROID开发
菜鸟系列之一:在 UBUNTU LINUX下搭建 ANDROID开发第一步:激活 ROOT用户.................................................... ...
aix oracle备份命令,Aix/Linux下自动备份oracle数据库
曾经有个同事,来回操作开发和生产的数据库,结果误删了生产的数据库,那种心情我想不是一般人能理解的,虽然说oracle可以有方法还原,但并不是彻底的. 所以,在工作中,不管是开发还是维护,备份数据库是非 ...
linux远程获取文件,Linux下实现获取远程机器文件
创建公钥秘钥实现无密码登录后即可获取到文件内容了!! A:xxx.xxx.6.xxx B:xxx.xxx.xxx.x 一.创建 A机器 ssh-keygen -t rsa 二.拷贝--将生成的公钥复制 ...

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

文章目录

步骤

第一步

第二步

第三步

第四步

第五步

总结

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

最新文章

热门文章