让系统自动选择空闲的GPU设备!帮你一次解决抢卡争端

项目地址:QuantumLiu / tf_gpu_manager

更新:支持pytorch

使用

git clone https://github.com/QuantumLiu/tf_gpu_manager

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

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

导入张量流为tf 从经理进口GPUManager 来自keras.layers LSTM GM = GPUManager() 与gm.auto_choice(): X = tf.placeholder(tf.float32,形状=(无,20,64)) Y = LSTM(32)(x)的

背景

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

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

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

  1. 一般来说,一个深度学习项目需要一个实验室或者小组的多人合作完成,要共享一台或几台工作站。一个主机多个GPU比较方便。
  2. 实验需要试多组参数或者对比试验。多GPU并行跑省时间。
  3. 模型计算量大,需要将模型不同分配在多个GPU上计算。

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

比如Tensorflow,在tensorflow \ python \ framework中定义了设备函数,返回一个用户执行操作的GPU设备的上下文管理器对象。

def device(device_name_or_function): “”“Wrapper for Graph.device()”使用默认图形。 看到 @ {} tf.Graph.device 更多细节。 ARGS: device_name_or_function:在上下文中使用的设备名称或功能。 返回: A context manager that specifies the default device to use for newly created ops. """ return get_default_graph().device(device_name_or_function)

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

with tf.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)

那么问题来了:

  1. 在写训练脚本时怎么知道哪个GPU是空闲可用的?
  2. 同组的人做实验和我冲突怎么办?
  3. 将来某个时刻运行这个脚本的时候是不是还要根据情况修改?
  4. 同行用我的代码复现实验,GPU配置环境不一样,他们甚至可能没有GPU,又要改代码?

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

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

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

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

实现

如何高效获取GPU状态信息

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

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

通过执行命令:

nvidia-smi --help-query-gpu

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

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

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

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

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=lambdav:(not'Not Support'inv)#lambda表达式,显卡是否滋瓷power management(笔记本可能不滋瓷) 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: 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]

如何衡量GPU空闲度

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

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

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

最后,我们用三个指标衡量:

  1. 显存剩余空间
  2. 显存剩余比例
  3. 当前功率/额定功率

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

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

如,按显存使用:

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)

完整脚本

我们定义一个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'

原文发布于微信公众号 - AI研习社(okweiwu)

原文发表时间:2017-09-04

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏python开发者

基于Python使用SVM识别简单的字符验证码的完整代码开源分享

关键字:Python,SVM,字符验证码,机器学习,验证码识别 1   概述 基于Python使用SVM识别简单的验证字符串的完整代码开源分享。 因为目前有了更...

43010
来自专栏企鹅号快讯

使用机器学习预测天气

作者:笨熊 本章是使用机器学习预测天气系列教程的第一部分,使用Python和机器学习来构建模型,根据从Weather Underground收集的数据来预测天气...

4425
来自专栏AI研习社

让 TensorFlow 估算器的推断提速百倍,我是怎么做到的?

TensorFlow 估算器提供了一套中阶 API 用于编写、训练与使用机器学习模型,尤其是深度学习模型。在这篇博文中,我们描述了如何通过使用异步执行来避免每次...

1552
来自专栏iOSDevLog

quickdraw_datasetQuick Draw!数据集

https://console.cloud.google.com/storage/browser/quickdraw_dataset

4912
来自专栏机器之心

作为TensorFlow的底层语言,你会用C++构建深度神经网络吗?

6369
来自专栏H2Cloud

A星路径搜索

摘要:   在人工智能中有一类问题是有确定解的,如路径、五子棋等,这样的问题非常适合使用搜索来解决。 路径搜索是一个很有趣的问题,在人工智能中算是很基础的问题。...

5314
来自专栏ATYUN订阅号

使用Java部署训练好的Keras深度学习模型

Keras库为深度学习提供了一个相对简单的接口,使神经网络可以被大众使用。然而,我们面临的挑战之一是将Keras的探索模型转化为产品模型。Keras是用Pyth...

6214
来自专栏AI科技评论

开发 | 如何为TensorFlow和PyTorch自动选择空闲GPU,解决抢卡争端

AI科技评论按:本文作者天清,原文载于其知乎专栏 世界那么大我想写代码,AI科技评论获授权发布。 项目地址:https://github.com/Quantum...

4058
来自专栏大数据挖掘DT机器学习

R语言关联规则可视化:扩展包arulesViz的介绍

关联规则挖掘是一种流行的数据挖掘方法,在R语言中为扩展包arules。然而,挖掘关联规则往往导致非常多的规则,使分析师需要通过查询所有的规则才能发现有趣的规则。...

5858
来自专栏大数据挖掘DT机器学习

R语言进行中文分词,并对6W条微博聚类

由于时间较紧,且人手不够,不能采用分类方法,主要是没有时间人工分类一部分生成训练集……所以只能用聚类方法,聚类最简单的方法无外乎:K-means与层次聚类。 尝...

3915

扫码关注云+社区

领取腾讯云代金券