雷鋒網(wǎng)按：本文為雷鋒字幕組編譯的技術(shù)博客，原標(biāo)題  Multithreaded predictions with TensorFlow Estimators，作者為 Archy de Berker 。

翻譯 | 李晶     校對 | 陳濤      整理 | MY   

TensorFlow 估算器提供了一套中階 API 用于編寫、訓(xùn)練與使用機(jī)器學(xué)習(xí)模型，尤其是深度學(xué)習(xí)模型。在這篇博文中，我們描述了如何通過使用異步執(zhí)行來避免每次調(diào)用預(yù)測方法時都需重載模型，從而讓 TF 估算器的推斷提速超過百倍。

什么是 TF 估算器？

TensorFlow 估算器于 2017 年年中被提出，首次出現(xiàn)在 KDD 的白皮書中。其設(shè)計目標(biāo)（如下面的兩分鐘視頻中所總結(jié)的）值得稱贊：將重復(fù)且容易出錯的任務(wù)自動化，將最佳實踐進(jìn)行封裝，保證了從訓(xùn)練到部署的順利執(zhí)行，所有這一切都以 scikit-learn 風(fēng)格進(jìn)行封裝。

2017 年 Martin Wicke 在介紹估算器接口。視頻來源：Google Developers, KDD 2017.

核心概念總結(jié)：用戶在 model_fn 中指定其模型中的關(guān)鍵點，使用條件語句來區(qū)分在訓(xùn)練和推斷中的不同操作。其中添加了一系列的 input_fns 來描述如何處理數(shù)據(jù)，可選擇為訓(xùn)練、評估和推斷分別指定各自的 input_fns 。

這些函數(shù)被 tf.estimator.Estimator 類調(diào)用并返回一個初始化的估算器。通過此估算器，可以調(diào)用 .train、.eval和 .predict 函數(shù)，而不用關(guān)心圖和會話，這兩個組件在基礎(chǔ)的 TensorFlow 設(shè)置中比較難用。

估算器接口。圖片來自 whitepaper （Cheng et al, 2017）

想獲得完整的實踐介紹，onfido blog 頁面提供了一個很棒的教程，該教程還包括 TensorFlow Dataset 和 Experiment 類（已棄用）。你可以在開始操作之前，先嘗試各種預(yù)先打包的估算器。

估算器面臨的挑戰(zhàn)

TensorFlow 是一個嵌合體：許多好的想法碰撞在一起，然而總體結(jié)構(gòu)并不完善。在這樣的背景下，估算器被提了出來，它需要與傳統(tǒng)的基于圖和會話的設(shè)計模式進(jìn)行競爭，而后者更為開發(fā)者所熟悉。開發(fā)者對估算器的接受也受到其代碼庫的混亂集成所影響，代碼庫中充滿了即將棄用的警告以及幾個明顯特征的遺漏（如 早期停止）。

因為其良好的默認(rèn)檢查點和 Tensorboard 集成，估算器在訓(xùn)練中使用起來很方便。然而，我們認(rèn)為推斷的接口有點不大直觀。

估算器的一個核心設(shè)計準(zhǔn)則是每次調(diào)用方法（.predict、.eval、.train）時都會重新對圖初始化。這不是很合理，下面所引用的原始論文對此進(jìn)行了總結(jié)：

為了確保封裝，每次調(diào)用方法時，估算器都會重新創(chuàng)建一個新圖，或許還會重載檢查點。重建圖的代價是很昂貴的，因而圖可以被緩存起來，從而減少在循環(huán)中執(zhí)行評估或預(yù)測的代價。但是，我們發(fā)現(xiàn)顯式重建圖還是很有用的，即使在明顯犧牲性能的情況下。

「TensorFlow 估算器：在高階機(jī)器學(xué)習(xí)框架下實現(xiàn)間接性和靈活性」，第 4 頁，作者  Cheng  等人

也就是說：在每次調(diào)用方法【train、predict、eval】時，都會重新構(gòu)建 TensorFlow 圖，并重新加載檢查點。要理解為什么會這樣，以及這會引起什么問題，我們需要深入了解這些方法的約定。

TF 估算器方法的約定

.train、.eval、.predict 都會用到 tensorflow 稱為 input_fn 的函數(shù)。調(diào)用此函數(shù)會返回一批數(shù)據(jù)。

通常由某種類型的生成器提供數(shù)據(jù)，這些生成器分批讀取數(shù)據(jù)，執(zhí)行預(yù)處理，并把它們傳遞給估算器。它們可以與 tf.Dataset 很好地結(jié)合在一起使用，tf.Dataset 能夠使上述過程（載入, 處理, 傳遞）并行化運行。

這意味著對于估算器而言，訓(xùn)練循環(huán)是在內(nèi)部進(jìn)行的。這樣做很有道理，正如白皮書中所強(qiáng)調(diào)的：

因為訓(xùn)練循環(huán)非常普遍，對其的最好實現(xiàn)應(yīng)該是移除許多重復(fù)的用戶代碼。這在理論上很簡單，我們可以避免由此產(chǎn)生的一些錯誤，不讓用戶為此而煩惱。因此，估算器實現(xiàn)并控制了訓(xùn)練循環(huán)。

「TensorFlow 估算器：在高階機(jī)器學(xué)習(xí)框架下實現(xiàn)間接性和靈活性」，第 5 頁，作者  Cheng  等人

這樣的設(shè)計可以很好地滿足需要預(yù)先對送入估算器的數(shù)據(jù)進(jìn)行指定的情況。該使用場景常出現(xiàn)在訓(xùn)練和評估中。

但是實際使用該模型進(jìn)行推斷的效果如何呢？

原始的推斷

假設(shè)我們想要將訓(xùn)練過的估算器用于另外一個任務(wù)，同樣是使用 Python。我們通常希望在一個工作流程中組合使用多個模型，例如使用語言模型作為自動語音轉(zhuǎn)錄或光學(xué)字符識別中定向搜索的補(bǔ)充。

為了簡化代碼庫，我們使用預(yù)打包的 Iris 數(shù)據(jù)集和估算器來模擬這種情況。假設(shè)我們有一種花卉推薦過程，它會不時地生成數(shù)據(jù)，并且每次都會從我們的估算器中讀取預(yù)測值。

每次生成推薦的候選時，該搜索過程都會調(diào)用我們的估算器。如果采用估算器的原始的實現(xiàn)方式，那么會非常緩慢，因為每次調(diào)用 flower_estimator.predict 都會重載估算器。

FlowerClassifier 類是對估算器的簡單包裝，它可能看起來像：

完整的代碼見 https://github.com/ElementAI/multithreaded-estimators/blob/1d0fba758d183193a822b8e44bda98a9443b456d/threaded_estimator/models.py#L12.

估算器的 .predict 方法已經(jīng)被封裝，所以調(diào)用 FlowerClassifier.predict() 會返回一個經(jīng)過訓(xùn)練的估算器的預(yù)測值。

完整代碼見 https://github.com/ElementAI/multithreaded-estimators/blob/master/threaded_estimator/tests/test_flower_estimator.py

但是現(xiàn)在每次我們想要分析一個新實例的時候，我們最終都會重新初始化整個模型！如果我們正在處理的任務(wù)代價很高，并且涉及到對模型的大量調(diào)用，那么效率就會嚴(yán)重下降。

緩存估算器來推斷

我們需要找到一種方法：僅調(diào)用一次 predict 方法，同時保證還能向生成器傳入新樣本。但是因為我們希望執(zhí)行其他中間計算，我們需要在單獨的線程中配置該生成器。

這是一個 生產(chǎn)者-消費者問題 的例子，在 Python 中可以使用隊列輕松解決。我們將使用兩個隊列以一種線程安全的方式移動數(shù)據(jù)，一個隊列用于保存輸入，另外一個隊列返回輸出：

乍看起來不大直觀，我們通過一個例子仔細(xì)研究一下到底發(fā)生了什么：

1. [主線程]： 用戶調(diào)用 .predict 方法

2. [主線程]： 將一系列新的數(shù)據(jù)被添加到 input_queue

3. [輔助線程]：數(shù)據(jù)生成器將從 input_queue 中生成一個輸入實例

4. [輔助線程]：該輸入實例被傳遞給模型

5. [輔助線程]：模型把生成的輸出實例添加到 output_queue

6. [主線程]: 調(diào)用封裝好的模型，返回 output_queue 中的最新項

在這個實現(xiàn)方案中，Python queues 的行為至關(guān)重要：如果隊列為空，則對 input_queue.get() 的調(diào)用會被先掛起，意味著生成器未被阻礙，只有數(shù)據(jù)被加入隊列后，才會繼續(xù)生成實例。

結(jié)果顯示整個會話過程中僅載入了一次模型。在 2017 款 MacBook Pro（沒有 GPU）的開發(fā)環(huán)境下運行，相比于原始實現(xiàn)，預(yù)測 100 個樣本類別的速度提升了大約 150 倍。

使用線程可能有些繁瑣，但是他們能把推斷的速度顯著加快。全部源代碼請見 https://github.com/ElementAI/multithreaded-estimators/blob/1d0fba758d183193a822b8e44bda98a9443b456d/threaded_estimator/models.py#L171.

需要注意的是，我們沒有對這個問題的其他解決方案進(jìn)行完全探索。我們可以使用 generator.send() 方法將實例注入數(shù)據(jù)生成器，我們也可以嘗試手動加載檢查點以執(zhí)行推理。我們發(fā)現(xiàn)這種特殊的方法非常有用，并且有很好的通用性，所以我們將其公之于眾：如果你發(fā)現(xiàn)這個問題還有其他的解決方案，我們愿聞其詳。

代碼

你可以在 Github 中找到代碼： https://github.com/ElementAI/multithreaded-estimators

我們提供了本文中討論到的類，一些測試和 Dockerfile，以幫助你啟動和運行環(huán)境。如果您覺得可以改進(jìn)代碼，隨時歡迎提交 Pull 請求。如果你更喜歡使用裝飾器，我們還有一個更復(fù)雜的版本，請參閱 decorator-refactor 分支。

感謝 Majid Laali 的原始想法和 Element AI 的整個 NLP 團(tuán)隊的編輯與建議。

想知道更多深度學(xué)習(xí)的技巧，訂閱  Element AI Lab Blog。

原文鏈接： https://medium.com/element-ai-research-lab/multithreaded-predictions-with-tensorflow-estimators-eb041861da07

雷鋒網(wǎng)雷鋒網(wǎng)

雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。