【解決搶卡爭端】為Tensorflow和PyTorch自動選擇空閑GPU

項目地址：QuantumLiu/tf_gpu_manager

***

更新：支持pytorch

***

使用

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

把manager.py放到你訓練的目錄就行。

直接使用with gm.auto_choice()自動選擇設備進行接下來代碼塊的操作。

import tensorflow as tfnfrom manager import GPUManagernfrom keras.layers LSTMngm=GPUManager()nwith gm.auto_choice():n x = tf.placeholder(tf.float32, shape=(None, 20, 64))n y = LSTM(32)(x) n

背景

隨著深度學習技術快速的發展，深度學習任務的數據和計算規模也越來越大，想要做出個像樣的work，沒有一台powerful的GPU工作站是萬萬不能的。

除了要求單卡性能強大，GPU數量多也很重要。

因為以下幾點原因，多GPU工作站已經成了各大實驗室的標配：

1. 一般來說，一個深度學習項目需要一個實驗室或者小組的多人合作完成，要共享一台或幾台工作站。一個host多個GPU比較方便。
2. 實驗需要試多組參數或者對比試驗。多GPU並行跑省時間。
3. 模型計算量大，需要將模型不同分配在多個GPU上計算。

現在，Tensorflow、pytorch等主流深度學習框架都支持多GPU訓練。

比如Tensorflow，在 tensorflowpythonframework 中定義了device函數，返回一個用來執行操作的GPU設備的context manager對象。

def device(device_name_or_function):n """Wrapper for `Graph.device()` using the default graph.nn Seen @{tf.Graph.device}n for more details.nn Args:n device_name_or_function: The device name or function to use inn the context.nn Returns:n A context manager that specifies the default device to use for newlyn created ops.n """n return get_default_graph().device(device_name_or_function)n

在我們的訓練腳本中使用with語句就可以指定接下來的操作在某個GPU上進行。

with tf.device(/gpu:2):n a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name=a)n b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name=b)n c = tf.matmul(a, b)n

那麼問題來了：

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-gpun

我們得到了所有支持的查詢參數（太多了不一一枚舉）

最有用的參數老司機給我們總結出來了：

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

於是我們有了基本思路，用os.popen執行相關命令，解析返迴文本信息。

def parse(line,qargs):n n line:n a line of textn qargs:n query argumentsn return:n a dict of gpu infosn Pasing a line of csv format text returned by nvidia-smin 解析一行nvidia-smi返回的csv格式文本n n numberic_args = [memory.free, memory.total, power.draw, power.limit]#可計數的參數n power_manage_enable=lambda vnot Not Support in v)#lambda表達式，顯卡是否滋瓷power management（筆記本可能不滋瓷）n to_numberic=lambda v:float(v.upper().strip().replace(MIB,).replace(W,))#帶單位字元串去掉單位n process = lambda k,v:((int(to_numberic(v)) if power_manage_enable(v) else 1) if k in numberic_args else v.strip())n return {k:process(k,v) for k,v in zip(qargs,line.strip().split(,))}nndef query_gpu(qargs=[]):n n qargs:n query argumentsn return:n a list of dictn Querying GPUs infosn 查詢GPU信息n n qargs =[index,gpu_name, memory.free, memory.total, power.draw, power.limit]+ qargsn cmd = nvidia-smi --query-gpu={} --format=csv,noheader.format(,.join(qargs))n results = os.popen(cmd).readlines()n return [parse(line,qargs) for line in results]n

如何衡量GPU空閑度

現在已經能獲取GPU狀態了，但是要怎麼衡量GPU空閑度並排序呢？

深度學習領域，GPU空閑度可以主要用兩個指標衡量：顯存空閑和功率空閑。

顯存佔用又分絕對空間佔用和佔用比例。

最後，我們用三個指標衡量：

1. 顯存剩餘空間
2. 顯存剩餘比例
3. 當前功率/額定功率

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

我們使用內置函數sorted來對可使用GPU進行排序。

如，按顯存使用：

def _sort_by_memory(self,gpus,by_size=False):n if by_size:n print(Sorted by free memory size)n return sorted(gpus,key=lambda d:d[memory.free],reverse=True)n else:n print(Sorted by free memory rate)n return sorted(gpus,key=lambda d:float(d[memory.free])/ d[memory.total],reverse=True)n

完整腳本

我們定義一個GPUManager類，在他的實例對象的存活周期里會更新GPU狀態、記錄已被分配的GPU。

實例化後，通過調用auto_choice方法直接返回一個tf.device對象。

同時，考慮到用戶計算機可能沒有GPU，加入異常處理機制。

def check_gpus():n n GPU available checkn reference : http://feisky.xyz/machine-learning/tensorflow/gpu_list.htmln n all_gpus = [x.name for x in device_lib.list_local_devices() if x.device_type == GPU]n if not all_gpus:n print(This script could only be used to manage NVIDIA GPUs,but no GPU found in your device)n return Falsen elif not NVIDIA System Management in os.popen(nvidia-smi -h).read():n print("nvidia-smi tool not found.")n return Falsen return Truennif check_gpus():n def parse(line,qargs):n n line:n a line of textn qargs:n query argumentsn return:n a dict of gpu infosn Pasing a line of csv format text returned by nvidia-smin 解析一行nvidia-smi返回的csv格式文本n n numberic_args = [memory.free, memory.total, power.draw, power.limit]#可計數的參數n power_manage_enable=lambda vnot Not Support in v)#lambda表達式，顯卡是否滋瓷power management（筆記本可能不滋瓷）n to_numberic=lambda v:float(v.upper().strip().replace(MIB,).replace(W,))#帶單位字元串去掉單位n process = lambda k,v:((int(to_numberic(v)) if power_manage_enable(v) else 1) if k in numberic_args else v.strip())n return {k:process(k,v) for k,v in zip(qargs,line.strip().split(,))}nn def query_gpu(qargs=[]):n n qargs:n query argumentsn return:n a list of dictn Querying GPUs infosn 查詢GPU信息n n qargs =[index,gpu_name, memory.free, memory.total, power.draw, power.limit]+ qargsn cmd = nvidia-smi --query-gpu={} --format=csv,noheader.format(,.join(qargs))n results = os.popen(cmd).readlines()n return [parse(line,qargs) for line in results]nn def by_power(d):n n helper function fo sorting gpus by powern n power_infos=(d[power.draw],d[power.limit])n if any(v==1 for v in power_infos):n print(Power management unable for GPU {}.format(d[index]))n return 1n return float(d[power.draw])/d[power.limit]nn class GPUManager():n n qargs:n query argumentsn A manager which can list all available GPU devicesn and sort them and choice the most free one.Unspecified n ones pref.n GPU設備管理器，考慮列舉出所有可用GPU設備，並加以排序，自動選出n 最空閑的設備。在一個GPUManager對象內會記錄每個GPU是否已被指定，n 優先選擇未指定的GPU。n n def __init__(self,qargs=[]):n n n self.qargs=qargsn self.gpus=query_gpu(qargs)n for gpu in self.gpus:n gpu[specified]=Falsen self.gpu_num=len(self.gpus)nn def _sort_by_memory(self,gpus,by_size=False):n if by_size:n print(Sorted by free memory size)n return sorted(gpus,key=lambda d:d[memory.free],reverse=True)n else:n print(Sorted by free memory rate)n return sorted(gpus,key=lambda d:float(d[memory.free])/ d[memory.total],reverse=True)nn def _sort_by_power(self,gpus):n return sorted(gpus,key=by_power)nn def _sort_by_custom(self,gpus,key,reverse=False,qargs=[]):n if isinstance(key,str) and (key in qargs):n return sorted(gpus,key=lambda d:d[key],reverse=reverse)n if isinstance(key,type(lambda a:a)):n return sorted(gpus,key=key,reverse=reverse)n raise ValueError("The argument key must be a function or a key in query args,please read the documention of nvidia-smi")nn def auto_choice(self,mode=0):n n mode:n 0:(default)sorted by free memory sizen return:n a TF device objectn Auto choice the freest GPU device,not specifiedn ones n 自動選擇最空閑GPUn n for old_infos,new_infos in zip(self.gpus,query_gpu(self.qargs)):n old_infos.update(new_infos)n unspecified_gpus=[gpu for gpu in self.gpus if not gpu[specified]] or self.gpusnn if mode==0:n print(Choosing the GPU device has largest free memory...)n chosen_gpu=self._sort_by_memory(unspecified_gpus,True)[0]n elif mode==1:n print(Choosing the GPU device has highest free memory rate...)n chosen_gpu=self._sort_by_power(unspecified_gpus)[0]n elif mode==2:n print(Choosing the GPU device by power...)n chosen_gpu=self._sort_by_power(unspecified_gpus)[0]n else:n print(Given an unaviliable mode,will be chosen by memory)n chosen_gpu=self._sort_by_memory(unspecified_gpus)[0]n chosen_gpu[specified]=Truen index=chosen_gpu[index]n print(Using GPU {i}:n{info}.format(i=index,info=n.join([str(k)+:+str(v) for k,v in chosen_gpu.items()])))n return tf.device(/gpu:{}.format(index))nelse:n raise ImportError(GPU available check failed)n

推薦閱讀：