雷鋒網按：本文作者天清， 原文載於其知乎專欄世界那麼大我想寫代碼， 雷鋒網獲其授權發佈。
項目位址：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工作站已經成了各大實驗室的標配：
一般來說， 一個深度學習專案需要一個實驗室或者小組的多人合作完成，
實驗需要試多組參數或者對比試驗。 多GPU並行跑省時間。
模型計算量大， 需要將模型不同分配在多個GPU上計算。
現在， Tensorflow、pytorch等主流深度學習框架都支持多GPU訓練。
比如Tensorflow， 在 tensorflowpython ramework中定義了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可以查詢顯卡資訊，
每次訓練前執行這個命令， 再與良好團隊保持良好的溝通可以解決上述1、2兩個問題， 但是3、4兩個問題還是不好解決。
而且經常和師兄弟、同事搶卡豈不是影響效率？
我們需要一種解決方案， 能夠實現不修改腳本、不需要和組員溝通，
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 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}:{info}'.format(i=index,info=''.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'