如何为TensorFlow和PyTorch自动选择空闲GPU解决抢卡争端_[#第一枪]

按：本文作者天清，原文载于其知乎专栏世界那么大我想写代码，雷锋网获其授权发布。

项目地址：QuantumLiu/tf_gpu_manager

***

更新：支持pytorch

***

使用

gitclonehttps://github.com/QuantumLiu/tf_gpu_manager

把manager.py放到你训练的目录就行。

直接使用with gm.auto_choice()自动选择设备进行接下来代码块的操作。

import tensorflow as tf

from manager import GPUManager

from keras.layers LSTM

gm=GPUManager()

with gm.auto_choice():

x=tf.placeholder(tf.float32,shape=(None,20,64))

y=LSTM(32)(x)

背景

随着深度学习技术快速的发展，深度学习任务的数据和计算规模也越来越大，想要做出个像样的work，没有一台powerful的GPU工作站是万万不能的。

除了要求单卡性能强大，GPU数量多也很重要。

因为以下几点原因，多GPU工作站已经成了各大实验室的标配：

一般来说，一个深度学习项目需要一个实验室或者小组的多人合作完成，要共享一台或几台工作站。一个host多个GPU比较方便。

实验需要试多组参数或者对比试验。多GPU并行跑省时间。

模型计算量大，需要将模型不同分配在多个GPU上计算。

现在，Tensorflow、pytorch等主流深度学习框架都支持多GPU训练。

比如Tensorflow，在 tensorflow\python\framework 中定义了device函数，返回一个用来执行操作的GPU设备的context manager对象。

def device(device_name_or_function):

"""Wrapperfor`Graph.device()`usingthedefaultgraph.

See

@{tf.Graph.device}

formoredetails.

Args:

device_name_or_function:Thedevicenameorfunctiontousein thecontext.

Returns:

Acontextmanagerthatspecifiesthedefaultdevicetousefornewly createdops.

"""

return get_default_graph().device(device_name_or_function)

在我们的训练脚本中使用with语句就可以指定接下来的操作在某个GPU上进行。

withtf.device('/gpu:2'):

a=tf.constant([1.0,2.0,3.0,4.0,5.0,6.0],shape=[2,3],name='a')

b=tf.constant([1.0,2.0,3.0,4.0,5.0,6.0],shape=[3,2],name='b')

c=tf.matmul(a,b)

那么问题来了：

在写训练脚本时怎么知道哪个GPU是空闲可用的？

同组的人做实验和我冲突怎么办？

将来某个时刻运行这个脚本的时候是不是还要根据情况修改？

同行用我的代码复现实验，GPU配置环境不一样，他们甚至可能没有GPU，又要改代码？

当然，上道儿的开发者都知道nvidia-smi可以查询显卡信息，查看GPU显存、温度、功率使用，然后选择合适的GPU。

<img src="https://static.leiphone.com/uploads/new/article/pic/201708/925d4b6331f0d6a55b25b9598bf86fc0.png" data-rawwidth="1109" data-rawheight="727" class="origin_image zh-lightbox-thumb" width="1109" data-original="https://pic4.zhimg.com/v2-4e889c4ab7942cc4056249f2f063f28f_r.png" _src="https://static.leiphone.com/uploads/new/article/pic/201708/925d4b6331f0d6a55b25b9598bf86fc0.png"/>

每次训练前执行这个命令，再与良好团队保持良好的沟通可以解决上述1、2两个问题，但是3、4两个问题还是不好解决。

而且经常和师兄弟、同事抢卡岂不是影响效率？

我们需要一种解决方案，能够实现不修改脚本、不需要和组员沟通，自动选择空闲GPU设备。

实现如何高效获取GPU状态信息

nvidia-smi是一个由NVIDIA官方提供的GPU状态管理、监控命令行软件。和其他命令行软件一样，nvidia-smi也有许多argument。

通过阅读文档，以及学习老司机的经验，我们知道--query-gpu这个option可以指定查询GPU状态信息，并返回格式化信息。

<img src="https://static.leiphone.com/uploads/new/article/pic/201708/3b2a2cee5d31afdcdeb65eac33ab0941.png" data-rawwidth="1465" data-rawheight="489" class="origin_image zh-lightbox-thumb" width="1465" data-original="https://pic4.zhimg.com/v2-b25096ab7d85f794c49a536738c21ff7_r.png" _src="https://static.leiphone.com/uploads/new/article/pic/201708/3b2a2cee5d31afdcdeb65eac33ab0941.png"/>

<img src="https://static.leiphone.com/uploads/new/article/pic/201708/9b5a148bfe1bc451cb58565fdeb1d3af.png" data-rawwidth="799" data-rawheight="275" class="origin_image zh-lightbox-thumb" width="799" data-original="https://pic2.zhimg.com/v2-f0f1924e0b3e10cdcb14ff759b82ef15_r.png" _src="https://static.leiphone.com/uploads/new/article/pic/201708/9b5a148bfe1bc451cb58565fdeb1d3af.png"/>

通过执行命令：

nvidia-smi--help-query-gpu

我们得到了所有支持的查询参数（太多了不一一枚举）

最有用的参数老司机给我们总结出来了：

<img src="https://static.leiphone.com/uploads/new/article/pic/201708/be63565f3373a5759a5d2e548c0adb6a.png" data-rawwidth="696" data-rawheight="185" class="origin_image zh-lightbox-thumb" width="696" data-original="https://pic2.zhimg.com/v2-db114bf08fa49676cb313d2688da7895_r.png" _src="https://static.leiphone.com/uploads/new/article/pic/201708/be63565f3373a5759a5d2e548c0adb6a.png"/>

还有我自己查到的index，name，power.draw, power.limit

<img src="https://static.leiphone.com/uploads/new/article/pic/201708/36ab0e6a17024f3892c1234934639e78.png" data-rawwidth="305" data-rawheight="185" class="content_image" width="305" _src="https://static.leiphone.com/uploads/new/article/pic/201708/36ab0e6a17024f3892c1234934639e78.png"/>

于是我们有了基本思路，用os.popen执行相关命令，解析返回文本信息。

def parse(line,qargs):

'''

line:

alineoftext

qargs:

queryarguments

return:

adictofgpuinfos

Pasingalineofcsvformattextreturnedbynvidia-smi

解析一行nvidia-smi返回的csv格式文本

'''

numberic_args=['memory.free','memory.total','power.draw','power.limit']#可计数的参数

power_manage_enable=lambdav:(not'NotSupport'inv)#lambda表达式，显卡是否滋瓷powermanagement（笔记本可能不滋瓷）

to_numberic=lambdav:float(v.upper().strip().replace('MIB','').replace('W',''))#带单位字符串去掉单位

process=lambdak,v:((int(to_numberic(v))ifpower_manage_enable(v)else1)ifkinnumberic_argselsev.strip())

return{k:process(k,v)fork,vinzip(qargs,line.strip().split(','))}

def query_gpu(qargs=[]):

'''

qargs:

queryarguments

return:

alistofdict

QueryingGPUsinfos

查询GPU信息

'''

qargs=['index','gpu_name','memory.free','memory.total','power.draw','power.limit']+qargs

cmd='nvidia-smi--query-gpu={}--format=csv,noheader'.format(','.join(qargs))

results=os.popen(cmd).readlines()

return [parse(line,qargs) for line in results]

如何衡量GPU空闲度

现在已经能获取GPU状态了，但是要怎么衡量GPU空闲度并排序呢？

深度学习领域，GPU空闲度可以主要用两个指标衡量：显存空闲和功率空闲。

显存占用又分绝对空间占用和占用比例。

最后，我们用三个指标衡量：

显存剩余空间

显存剩余比例

当前功率/额定功率

在之前，我们已经把所有GPU的信息存成了一个list，每个list是gpu信息的字典。

我们使用内置函数sorted来对可使用GPU进行排序。

如，按显存使用：

def_sort_by_memory(self,gpus,by_size=False):

if by_size:

print('Sortedbyfreememorysize')

return sorted(gpus,key=lambda d:d['memory.free'],reverse=True)

else:

print('Sortedbyfreememoryrate')

return sorted(gpus,key=lambda d:float(d['memory.free']) / d['memory.total'],reverse=True)

完整脚本

我们定义一个GPUManager类，在他的实例对象的存活周期里会更新GPU状态、记录已被分配的GPU。

实例化后，通过调用auto_choice方法直接返回一个tf.device对象。

同时，考虑到用户计算机可能没有GPU，加入异常处理机制。

def check_gpus():

'''

GPU available check

reference : http://feisky.xyz/machine-learning/tensorflow/gpu_list.html

'''

all_gpus = [x.name for x in device_lib.list_local_devices() if x.device_type == 'GPU']

if not all_gpus:

print('This script could only be used to manage NVIDIA GPUs,but no GPU found in your device')

return False

elif not 'NVIDIA System Management' in os.popen('nvidia-smi -h').read():

print("'nvidia-smi' tool not found.")

return False

return True

if check_gpus():

def parse(line,qargs):

'''

line:

a line of text

qargs:

query arguments

return:

a dict of gpu infos

Pasing a line of csv format text returned by nvidia-smi

解析一行nvidia-smi返回的csv格式文本

'''

numberic_args = ['memory.free', 'memory.total', 'power.draw', 'power.limit']#可计数的参数

power_manage_enable=lambda v:(not 'Not Support' in v)#lambda表达式，显卡是否滋瓷power management（笔记本可能不滋瓷）

to_numberic=lambda v:float(v.upper().strip().replace('MIB','').replace('W',''))#带单位字符串去掉单位

process = lambda k,v:((int(to_numberic(v)) if power_manage_enable(v) else 1) if k in numberic_args else v.strip())

return {k:process(k,v) for k,v in zip(qargs,line.strip().split(','))}

def query_gpu(qargs=[]) :

'''

qargs:

query arguments

return:

a list of dict

Querying GPUs infos

查询GPU信息

'''

qargs =['index','gpu_name', 'memory.free', 'memory.total', 'power.draw', 'power.limit']+ qargs

cmd = 'nvidia-smi --query-gpu={} --format=csv,noheader'.format(','.join(qargs))

results = os.popen(cmd).readlines()

return [parse(line,qargs) for line in results]

def by_power(d):

'''

helper function fo sorting gpus by power

'''

power_infos=(d['power.draw'],d['power.limit'])

if any(v==1 for v in power_infos):

print('Power management unable for GPU {}'.format(d['index']))

return 1

return float(d['power.draw'])/d['power.limit']

class GPUManager():

'''

qargs:

query arguments

A manager which can list all available GPU devices and sort them and choice the most free one.Unspecified ones pref.

GPU设备管理器，考虑列举出所有可用GPU设备，并加以排序，自动选出最空闲的设备。在一个GPUManager对象内会记录每个GPU是否已被指定，优先选择未指定的GPU。

'''

def __init__(self,qargs=[]):

'''

'''

self.qargs=qargs

self.gpus=query_gpu(qargs)

for gpu in self.gpus:

gpu['specified']=False

self.gpu_num=len(self.gpus)

def _sort_by_memory(self,gpus,by_size=False):

if by_size:

print('Sorted by free memory size')

return sorted(gpus,key=lambda d:d['memory.free'],reverse=True)

else:

print('Sorted by free memory rate')

return sorted(gpus,key=lambda d:float(d['memory.free'])/ d['memory.total'],reverse=True)

def _sort_by_power(self,gpus):

return sorted(gpus,key=by_power)

def _sort_by_custom(self,gpus,key,reverse=False,qargs=[]):

if isinstance(key,str) and (key in qargs):

return sorted(gpus,key=lambda d:d[key],reverse=reverse)

if isinstance(key,type(lambda a:a)):

return sorted(gpus,key=key,reverse=reverse)

raise ValueError("The argument 'key' must be a function or a key in query args,please read the documention of nvidia-smi")

def auto_choice(self,mode=0):

'''

mode:

0:(default)sorted by free memory size

return:

a TF device object

Auto choice the freest GPU device,not specified

ones

自动选择最空闲GPU

'''

for old_infos,new_infos in zip(self.gpus,query_gpu(self.qargs)):

old_infos.update(new_infos)

unspecified_gpus=[gpu for gpu in self.gpus if not gpu['specified']] or self.gpus

if mode==0:

print('Choosing the GPU device has largest free memory...')

chosen_gpu=self._sort_by_memory(unspecified_gpus,True)[0]

elif mode==1:

print('Choosing the GPU device has highest free memory rate...')

chosen_gpu=self._sort_by_power(unspecified_gpus)[0]

elif mode==2:

print('Choosing the GPU device by power...')

chosen_gpu=self._sort_by_power(unspecified_gpus)[0]

else:

print('Given an unaviliable mode,will be chosen by memory')

chosen_gpu=self._sort_by_memory(unspecified_gpus)[0]

chosen_gpu['specified']=True

index=chosen_gpu['index']

print('Using GPU {i}:\n{info}'.format(i=index,info='\n'.join([str(k)+':'+str(v) for k,v in chosen_gpu.items()])))

return tf.device('/gpu:{}'.format(index))

else:

raise ImportError('GPU available check failed'

雷锋网版权文章，未经授权禁止转载。详情见转载须知。

复写纸

钢制暖气片价格

果树苗木批发

开袋机价格