利用 Python 自动获取服务器中空闲 GPU

服务器资源紧张，每次要使用 GPU 的时候，都得先用命令 nvidia-smi 看一下 GPU 的使用情况，然后手工选择一个比较合适的 GPU 进行运算，本文通过 Python 实现了自动获取服务器中相对比较空闲的 GPU。但对于实际使用来说，这个代码好像也并有什么卵用，除非自己首先是知道了自己的模型需要占用多少显存，然后有针对性地进行选择。

运行环境

• Ubuntu 16.04
• Python 3.6.1

代码

import os
import numpy as np

# 获取每个 GPU 的剩余显存数，并存放到 tmp 文件中
os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
memory_gpu = [int(x.split()[2]) for x in open('tmp', 'r').readlines()]
print("剩余显存最大的 GPU 编号是： "np.argmax(memory_gpu))  # 获取剩余显存最多的 GPU 的编号
os.system('rm tmp')  # 删除临时生成的 tmp 文件

报错信息如下

sh: 1: cannot create tmp: Permission denied

原因：在当前文件夹下，该用户无对应写操作的权限。重新在有权限的区域进行运行即可。

获取到的每张 GPU 的空闲 memory 如下，临时存放于 tmp 文件中。

在 PyTorch 中自动选择 GPU

import os
import numpy as np
import torch

# 获取每个 GPU 的剩余显存数，并存放到 tmp 文件中
os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
memory_gpu = [int(x.split()[2]) for x in open('tmp', 'r').readlines()]
torch.cuda.set_device(np.argmax(memory_gpu)) 
os.system('rm tmp')  # 删除临时生成的 tmp 文件

参考文献

• python实现自动查找并使用空闲GPU