前言
首先,若是你如今已經很熟悉tf.data+estimator了,能夠把文章x掉了╮( ̄▽ ̄」」)╭python
可是!若是如今仍是在進行session.run(…)的話!尤爲是苦惱於GPU顯存都塞滿了利用率卻上不去的童鞋,這篇文章或許能夠給你打開新世界的大門噢( ̄∇ ̄)編程
若是發現通過一系列改良後訓練效率大大提升了,記得回來給小夕發小紅包( ̄∇ ̄)api
不過,這並非一篇怒貼一堆代碼,言(三)簡(言)意(兩)賅(語)就結束的CSDN文風的文章。。。因此伸手黨們也能夠X掉了╮( ̄▽ ̄」」)╭session
緣起
很早很早以前,在小夕剛接觸tensorflow和使用GPU加速計算的時候,就產生過一個疑惑。爲何顯卡的顯存都快滿了,GPU利用率還顯示這麼低呢?好浪費呀,可是又迫不得已。當時GPU利用率100%的狀況基本是僅存於一塊顯卡塞四、5個不費顯存的小任務的狀況。多線程
大部分狀況下寫出來的代碼train起來後是這樣的:app
能夠看到,雖然顯卡的顯存都塞滿了,可是顯卡功率(最左邊那一欄,114W和69W)和利用率(最右邊那一欄,35%和38%)卻遠遠沒有達到極限。大部分人的想法是,算了算了這不重要,我去作實驗了再見【wei笑】函數式編程
然而!若是你在作大型實驗,train一次跑幾天呢?這個細節會極大的影響你的實驗效率和DDL到來前的實驗次數!想一下,徹底同樣的model和設置,你的代碼要train一週,然而隔壁老王只須要train三天╮( ̄▽ ̄」」)╭函數
路人甲:我有256張顯卡
小夕:好了這篇文章你能夠X掉了
那麼,咱們有沒有可能一直這樣呢:
oop
是否是這功率和利用率看起來難以想象!不要懷疑這是PS的圖!這只是小夕的平常截圖!tricks用的好GPU利用率掉不下來99%,然鵝代碼寫的足夠蠢,也能夠上不去5%!post
那麼問題來了,究竟是什麼致使的這個差別呢?
不要急,咱們來放大一下那些gpu利用率只有30%幾的代碼在訓練時的gpu利用率的變化狀況(好像句子有點長
watch -n 0.1 nvidia-smi
ps:(可能掉幀太嚴重了看着不連貫╮( ̄▽ ̄"")╭,建議在本身的機器上試一下,會直觀的多~)
看!是否是一會兒就發現問題啦?能夠看到,其實gpu利用率並非一直在比較低的水平,而是頗有規律的週期性的從0漲到接近100再跌到0,再從新漲到100再跌回0。若是同時開着打印日誌的窗口,你就會發現這個週期剛好跟每一個訓練step的時長一致!也就是說,在每一個step,其實有一些時間並無花在GPU裏,那固然就是花在cpu裏啦。
那在cpu裏幹什麼的呢?固然就是load下一個batch、預處理這個batch以及在gpu上跑出結果後打印日誌、後處理、寫summary甚至保存模型等,這一系列的花銷都要靠cpu去完成。回顧一下咱們常寫的代碼:
create_graph() create_model_saver() create_summary_writer() create_session() do_init() for i in range(num_train_steps): load_batch(...) # cpu preprocess(...) # cpu feed_dict = {...} # cpu fetch_list = [...] # cpu buf = session.run(fetch_list, feed_dict) # gpu postprocess(buf) # cpu print(...) # cpu if i % x == 0: summary_writer.write(...) # cpu if i % xx == 0: model_saver.save(...) # cpu
看,尤爲是preprocess(…)任務比較重的話就容易致使代碼在cpu裏也要跑好一段時間,gpu利用率天然就會上不去並且呈現週期性變化啦。
那麼有沒有什麼辦法下降cpu時間,提升gpu時間呢?
一個很自(愚)然(蠢)的想法就是把一切訓練代碼都用tf的api重寫不就好啦,甚至最外層的那個for i in range(num_train_steps)其實均可以用tf.while_loop重寫呀。嗯,小夕還真的這麼嘗試過,而後發現
TF api這特喵的都是些什麼鬼!各類跟numpy和python內置函數重名卻行爲不一致是什麼鬼!臥槽這個api少了個參數我該怎麼辦?python裏就一行代碼就能搞定的事情我爲何寫了幾十行??
因此除了函數式編程的大牛,小夕極力的不建議重蹈覆轍!尤爲是咱們這些遇到彙編會哭,看到Lisp會崩潰的90後小仙女!
因此沒辦法把整個train loop都描述進計算圖了?
別怕別怕,好在後來其實tensorflow已經封裝了一個特別好(多)用(坑)的上層API來把整個train loop都能輕鬆的封裝在計算圖中,從而實現超級高的GPU利用率和訓練效率!
Estimator
不用管它爲啥叫Estimator,只須要知道,它把咱們剛纔想作的事情基本都給封裝好了就行。把剛纔的那個經典的寫法搬過來
1. create_model() 2. create_model_saver() 3. create_summary_writer() 4. create_session() 5. do_init() 6. for i in range(num_train_steps): 7. load_batch(...) # cpu 8. preprocess(...) # cpu 9. feed_dict = {...} # cpu 10. fetch_list = [...] # cpu 11. buf = session.run(fetch_list, feed_dict) # gpu 12. postprocess(buf) # cpu 13. print(...) # cpu 14. if i % x == 0: 15. summary_writer.write(...) # cpu 16. if i % xx == 0: 17. model_saver.save(...) # cpu
1-5行在estimator中都封裝好啦,你只須要把相關配置塞進estimator的RunConfig就能夠啦~
7-9行也封裝好啦,你只須要把數據集載入和預處理的相關代碼的函數塞給estimator.train的input_fn~
第10行也封裝好啦,你只須要把要fetch的loss、train_op丟進estimator的EstimatorSpec~
第11行也封裝好啦,你只須要把描述模型計算圖的函數塞給estimator的model_fn~
第12-13行不用操心細節了,global_step和loss自動完成了,剩下的丟給tf.Print和LoggingTensorHook吧~
第14-17行不用你寫了,自動完成了
╮(╯▽╰)╭
通過這麼一頓折騰,咱們發現GPU利用率大大提升啦~直逼80%甚至90%。那麼還有沒有能夠壓榨的空間呢?
其實這時仔細一分析就會發現雖然estimator把大部分的代碼寫進計算圖裏了,可是從數據的載入和預處理依然是在cpu裏串行進行呀,並且好比一個batch有128個樣本,那麼estimaor內部在run每一個step的時候仍是要等着這128個樣本串行的處理完才行。這顯然就是最後的瓶頸啦!有沒有辦法消除掉呢?·固然有,那就是
tf.data
TF的dataset API能夠說讓人又愛又恨了,它確實看似提供了一種把整個預處理都搬進計算圖進行並行化處理的途徑,可是!若是你真的徹底用tensorflow API來作複雜的預處理的話,真的會讓人瘋掉的QAQ所以,這裏在用tf.data以前,小夕極力的建議先把數據集儘量的transform成預處理後的樣子,包括作分詞、作截斷、作word2id等,不過padding和input_mask能夠留在TF裏面作,畢竟都只須要一行。
那作完這些預處理後,數據該怎麼存儲會更方便後續的讀取和處理呢?最最最建議的方式仍是使用tf.records來存儲,磁盤、內存的存儲和IO效率都會相比傳統方式更快一些,x和y也不用分開了。固然這樣的惟一的壞處就是不能直接打開看數據集╮( ̄▽ ̄」」)╭畢竟數據集被作成了二進制文件。
可是實在比較懶不想用tf.record的話,那麼小夕極力建議把x和y分開存儲,而且儘可能讓tf.data在讀取數據的時候作完上面的那些必要的預處理,以避開難用的字符串基礎操做API而且減輕訓練時的cpu和內存壓力。
tf.data還有一個很大的好處就是能夠很自然的支持以streaming的方式讀取數據,這樣在面對大數據集時就不會發生數據load完後發現顯卡被佔的尷尬事件了╮( ̄▽ ̄」」)╭
好像講了這麼久,仍是沒講怎麼用tf.data加速QAQ,來來來進入正題啦。
想一想哈,沒用tf.data的時候,咱們寫出來的代碼實際跑起來就是這個樣子的:
這也是文章開頭小夕解釋的爲何gpu利用率上不去而且週期性變化的重要緣由。那麼咱們能夠不能夠消除idle,像下面這樣讓prepare和train的過程並行進行呢?
固然能夠!那就是
prefetch
從prefetch的意思就能夠理解,那就是預先獲取下一個step要load的batch。使用tf.data裏面的叫作prefetch的神奇api就能夠輕鬆完成啦,這個api裏的參數buffer_size就是講的是額外的fetch多少份,好比buffer_size=1,而後咱們要prefetch的是batch的話,那麼模型每次prepare完一個batch後,就會自動再額外的prepare一個batch,這樣下一個train step到來的時候就能夠直接從內存中取走這個事先prepare好的batch啦。(詳情見後面)
等下,看上圖的話,有木有發現,若是prepare一個batch耗時很短的話確實兩全齊美,可是若是耗時比較久,尤爲一會兒prefetch好幾個batch的話,一旦prepare的用時超過了train一個step的用時,那麼每一個train step的性能就會受限於prepare的效率啦。放大一下這個問題的話以下圖所示
看,prepare用時過久反而會致使train完一個step後gpu空閒了(雖然其實下個step的batch可能已經prepare好了)
那麼能不能確保prepare階段的用時小於train階段的用時呢?
parallel mapping
一個很簡單的想法固然就是讓樣本並行處理啦~若是batch size是128,prefetch size=1,那麼準備一個batch要串行的跑128*2=256次的預處理,可是若是咱們開4個線程去跑,是否是就看起來快多啦。幸運的是咱們也不用本身手擼多線程了,tf.data.Dataset在map(預處理)函數裏有一個參數num_parallel_calls,給這個參數賦值就能夠並行parse啦。如圖,
這樣的話只要prefetch的buffer_size和map的num_parrellel_calls取得合適,基本就能夠實現不間斷的train啦,也就是幾乎達到100%的GPU利用率!
好啦,思想明白了,代碼就容易理解啦。不使用tf.record,直接從預處理好的純文本格式的數據集load數據時的典型過程以下
def build_input(..): x = tf.data.XXDataset(..) x = x.map(..., num_parallel_calls=N) # parellel y = tf.data.XXDataset(..) y = y.map(..., num_parallel_calls=N) dataset = tf.data.Dataset.zip((x, y)) dataset = dataset.repeat(num_epochs) if is_train: dataset = dataset.shuffle(..) dataset = dataset.batch(batch_size) dataset = dataset.prefetch(buffer_size=1) # prefetch iterator = dataset.make_xx_iterator() return iterator.get_next()
固然,若是用上tf.record後,就不用分別從x和y倆文件中讀數據啦,感興趣的童鞋可自行去了解一下。