0
| 本文作者： AI研習(xí)社-譯站 | 2018-07-19 18:21 |
雷鋒網(wǎng)按：本文為雷鋒字幕組編譯的技術(shù)博客，原標(biāo)題 Multithreaded predictions with TensorFlow Estimators，作者為 Archy de Berker 。
翻譯 | 李晶 校對(duì) | 陳濤 整理 | MY
TensorFlow 估算器提供了一套中階 API 用于編寫、訓(xùn)練與使用機(jī)器學(xué)習(xí)模型，尤其是深度學(xué)習(xí)模型。在這篇博文中，我們描述了如何通過使用異步執(zhí)行來避免每次調(diào)用預(yù)測(cè)方法時(shí)都需重載模型，從而讓 TF 估算器的推斷提速超過百倍。
什么是 TF 估算器？
TensorFlow 估算器于 2017 年年中被提出，首次出現(xiàn)在 KDD 的白皮書中。其設(shè)計(jì)目標(biāo)（如下面的兩分鐘視頻中所總結(jié)的）值得稱贊：將重復(fù)且容易出錯(cuò)的任務(wù)自動(dòng)化，將最佳實(shí)踐進(jìn)行封裝，保證了從訓(xùn)練到部署的順利執(zhí)行，所有這一切都以 scikit-learn 風(fēng)格進(jìn)行封裝。
2017 年 Martin Wicke 在介紹估算器接口。視頻來源：Google Developers, KDD 2017.
核心概念總結(jié)：用戶在 model_fn 中指定其模型中的關(guān)鍵點(diǎn)，使用條件語句來區(qū)分在訓(xùn)練和推斷中的不同操作。其中添加了一系列的 input_fns 來描述如何處理數(shù)據(jù)，可選擇為訓(xùn)練、評(píng)估和推斷分別指定各自的 input_fns 。
這些函數(shù)被 tf.estimator.Estimator 類調(diào)用并返回一個(gè)初始化的估算器。通過此估算器，可以調(diào)用 .train、.eval和 .predict 函數(shù)，而不用關(guān)心圖和會(huì)話，這兩個(gè)組件在基礎(chǔ)的 TensorFlow 設(shè)置中比較難用。
估算器接口。圖片來自 whitepaper （Cheng et al, 2017）
想獲得完整的實(shí)踐介紹，onfido blog 頁面提供了一個(gè)很棒的教程，該教程還包括 TensorFlow Dataset 和 Experiment 類（已棄用）。你可以在開始操作之前，先嘗試各種預(yù)先打包的估算器。
估算器面臨的挑戰(zhàn)
TensorFlow 是一個(gè)嵌合體：許多好的想法碰撞在一起，然而總體結(jié)構(gòu)并不完善。在這樣的背景下，估算器被提了出來，它需要與傳統(tǒng)的基于圖和會(huì)話的設(shè)計(jì)模式進(jìn)行競(jìng)爭(zhēng)，而后者更為開發(fā)者所熟悉。開發(fā)者對(duì)估算器的接受也受到其代碼庫(kù)的混亂集成所影響，代碼庫(kù)中充滿了即將棄用的警告以及幾個(gè)明顯特征的遺漏（如 早期停止）。
因?yàn)槠淞己玫哪J(rèn)檢查點(diǎn)和 Tensorboard 集成，估算器在訓(xùn)練中使用起來很方便。然而，我們認(rèn)為推斷的接口有點(diǎn)不大直觀。
估算器的一個(gè)核心設(shè)計(jì)準(zhǔn)則是每次調(diào)用方法（.predict、.eval、.train）時(shí)都會(huì)重新對(duì)圖初始化。這不是很合理，下面所引用的原始論文對(duì)此進(jìn)行了總結(jié)：
為了確保封裝，每次調(diào)用方法時(shí)，估算器都會(huì)重新創(chuàng)建一個(gè)新圖，或許還會(huì)重載檢查點(diǎn)。重建圖的代價(jià)是很昂貴的，因而圖可以被緩存起來，從而減少在循環(huán)中執(zhí)行評(píng)估或預(yù)測(cè)的代價(jià)。但是，我們發(fā)現(xiàn)顯式重建圖還是很有用的，即使在明顯犧牲性能的情況下。
「TensorFlow 估算器：在高階機(jī)器學(xué)習(xí)框架下實(shí)現(xiàn)間接性和靈活性」，第 4 頁，作者 Cheng 等人
也就是說：在每次調(diào)用方法【train、predict、eval】時(shí)，都會(huì)重新構(gòu)建 TensorFlow 圖，并重新加載檢查點(diǎn)。要理解為什么會(huì)這樣，以及這會(huì)引起什么問題，我們需要深入了解這些方法的約定。
TF 估算器方法的約定
.train、.eval、.predict 都會(huì)用到 tensorflow 稱為 input_fn 的函數(shù)。調(diào)用此函數(shù)會(huì)返回一批數(shù)據(jù)。
通常由某種類型的生成器提供數(shù)據(jù)，這些生成器分批讀取數(shù)據(jù)，執(zhí)行預(yù)處理，并把它們傳遞給估算器。它們可以與 tf.Dataset 很好地結(jié)合在一起使用，tf.Dataset 能夠使上述過程（載入, 處理, 傳遞）并行化運(yùn)行。
這意味著對(duì)于估算器而言，訓(xùn)練循環(huán)是在內(nèi)部進(jìn)行的。這樣做很有道理，正如白皮書中所強(qiáng)調(diào)的：
因?yàn)橛?xùn)練循環(huán)非常普遍，對(duì)其的最好實(shí)現(xiàn)應(yīng)該是移除許多重復(fù)的用戶代碼。這在理論上很簡(jiǎn)單，我們可以避免由此產(chǎn)生的一些錯(cuò)誤，不讓用戶為此而煩惱。因此，估算器實(shí)現(xiàn)并控制了訓(xùn)練循環(huán)。
「TensorFlow 估算器：在高階機(jī)器學(xué)習(xí)框架下實(shí)現(xiàn)間接性和靈活性」，第 5 頁，作者 Cheng 等人
這樣的設(shè)計(jì)可以很好地滿足需要預(yù)先對(duì)送入估算器的數(shù)據(jù)進(jìn)行指定的情況。該使用場(chǎng)景常出現(xiàn)在訓(xùn)練和評(píng)估中。
但是實(shí)際使用該模型進(jìn)行推斷的效果如何呢？
原始的推斷
假設(shè)我們想要將訓(xùn)練過的估算器用于另外一個(gè)任務(wù)，同樣是使用 Python。我們通常希望在一個(gè)工作流程中組合使用多個(gè)模型，例如使用語言模型作為自動(dòng)語音轉(zhuǎn)錄或光學(xué)字符識(shí)別中定向搜索的補(bǔ)充。
為了簡(jiǎn)化代碼庫(kù)，我們使用預(yù)打包的 Iris 數(shù)據(jù)集和估算器來模擬這種情況。假設(shè)我們有一種花卉推薦過程，它會(huì)不時(shí)地生成數(shù)據(jù)，并且每次都會(huì)從我們的估算器中讀取預(yù)測(cè)值。
每次生成推薦的候選時(shí)，該搜索過程都會(huì)調(diào)用我們的估算器。如果采用估算器的原始的實(shí)現(xiàn)方式，那么會(huì)非常緩慢，因?yàn)槊看握{(diào)用 flower_estimator.predict 都會(huì)重載估算器。
FlowerClassifier 類是對(duì)估算器的簡(jiǎn)單包裝，它可能看起來像：
估算器的 .predict 方法已經(jīng)被封裝，所以調(diào)用 FlowerClassifier.predict() 會(huì)返回一個(gè)經(jīng)過訓(xùn)練的估算器的預(yù)測(cè)值。
但是現(xiàn)在每次我們想要分析一個(gè)新實(shí)例的時(shí)候，我們最終都會(huì)重新初始化整個(gè)模型！如果我們正在處理的任務(wù)代價(jià)很高，并且涉及到對(duì)模型的大量調(diào)用，那么效率就會(huì)嚴(yán)重下降。
緩存估算器來推斷
我們需要找到一種方法：僅調(diào)用一次 predict 方法，同時(shí)保證還能向生成器傳入新樣本。但是因?yàn)槲覀兿Ｍ麍?zhí)行其他中間計(jì)算，我們需要在單獨(dú)的線程中配置該生成器。
這是一個(gè) 生產(chǎn)者-消費(fèi)者問題 的例子，在 Python 中可以使用隊(duì)列輕松解決。我們將使用兩個(gè)隊(duì)列以一種線程安全的方式移動(dòng)數(shù)據(jù)，一個(gè)隊(duì)列用于保存輸入，另外一個(gè)隊(duì)列返回輸出：
乍看起來不大直觀，我們通過一個(gè)例子仔細(xì)研究一下到底發(fā)生了什么：
[主線程]： 用戶調(diào)用 .predict 方法
[主線程]： 將一系列新的數(shù)據(jù)被添加到 input_queue
[輔助線程]：數(shù)據(jù)生成器將從 input_queue 中生成一個(gè)輸入實(shí)例
[輔助線程]：該輸入實(shí)例被傳遞給模型
[輔助線程]：模型把生成的輸出實(shí)例添加到 output_queue
[主線程]: 調(diào)用封裝好的模型，返回 output_queue 中的最新項(xiàng)
在這個(gè)實(shí)現(xiàn)方案中，Python queues 的行為至關(guān)重要：如果隊(duì)列為空，則對(duì) input_queue.get() 的調(diào)用會(huì)被先掛起，意味著生成器未被阻礙，只有數(shù)據(jù)被加入隊(duì)列后，才會(huì)繼續(xù)生成實(shí)例。
結(jié)果顯示整個(gè)會(huì)話過程中僅載入了一次模型。在 2017 款 MacBook Pro（沒有 GPU）的開發(fā)環(huán)境下運(yùn)行，相比于原始實(shí)現(xiàn)，預(yù)測(cè) 100 個(gè)樣本類別的速度提升了大約 150 倍。
使用線程可能有些繁瑣，但是他們能把推斷的速度顯著加快。全部源代碼請(qǐng)見 https://github.com/ElementAI/multithreaded-estimators/blob/1d0fba758d183193a822b8e44bda98a9443b456d/threaded_estimator/models.py#L171.
需要注意的是，我們沒有對(duì)這個(gè)問題的其他解決方案進(jìn)行完全探索。我們可以使用 generator.send() 方法將實(shí)例注入數(shù)據(jù)生成器，我們也可以嘗試手動(dòng)加載檢查點(diǎn)以執(zhí)行推理。我們發(fā)現(xiàn)這種特殊的方法非常有用，并且有很好的通用性，所以我們將其公之于眾：如果你發(fā)現(xiàn)這個(gè)問題還有其他的解決方案，我們?cè)嘎勂湓敗?/p>
代碼
你可以在 Github 中找到代碼： https://github.com/ElementAI/multithreaded-estimators
我們提供了本文中討論到的類，一些測(cè)試和 Dockerfile，以幫助你啟動(dòng)和運(yùn)行環(huán)境。如果您覺得可以改進(jìn)代碼，隨時(shí)歡迎提交 Pull 請(qǐng)求。如果你更喜歡使用裝飾器，我們還有一個(gè)更復(fù)雜的版本，請(qǐng)參閱 decorator-refactor 分支。
感謝 Majid Laali 的原始想法和 Element AI 的整個(gè) NLP 團(tuán)隊(duì)的編輯與建議。
想知道更多深度學(xué)習(xí)的技巧，訂閱 Element AI Lab Blog。
雷鋒網(wǎng)雷鋒網(wǎng)
雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
