雷锋网按:本文作者天清,原文载于其知乎专栏世界那么大我想写代码,雷锋网获其授权发布。
项目地址:QuantumLiu/tf_gpu_manager
***
更新:支持pytorch
***
git clone https://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):
"""Wrapper for `Graph.device()` using the default graph.
See
@{tf.Graph.device}
for more details.
Args:
device_name_or_function: The device name or function to use in the context.
Returns:
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)
那么问题来了:
在写训练脚本时怎么知道哪个GPU是空闲可用的?
同组的人做实验和我冲突怎么办?
将来某个时刻运行这个脚本的时候是不是还要根据情况修改?
同行用我的代码复现实验,GPU配置环境不一样,他们甚至可能没有GPU,又要改代码?
当然,上道儿的开发者都知道nvidia-smi可以查询显卡信息,查看GPU显存、温度、功率使用,然后选择合适的GPU。
每次训练前执行这个命令,再与良好团队保持良好的沟通可以解决上述1、2两个问题,但是3、4两个问题还是不好解决。
而且经常和师兄弟、同事抢卡岂不是影响效率?
我们需要一种解决方案,能够实现不修改脚本、不需要和组员沟通,自动选择空闲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的信息存成了一个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 Trueif 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'