Tensorflow 估算器的加速站，你捨得錯過嗎？

雷鋒網按：本文為雷鋒字幕組編譯的技術博客，原標題  Multithreaded predictions with TensorFlow Estimators，作者為 Archy de Berker 。

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

TensorFlow 估算器提供了一套中階 API 用於編寫、訓練與使用機器學習模型，尤其是深度學習模型。在這篇博文中，我們描述了如何通過使用非同步執行來避免每次調用預測方法時都需重載模型，從而讓 TF 估算器的推斷提速超過百倍。

什麼是 TF 估算器？

TensorFlow 估算器於 2017 年年中被提出，首次出現在 KDD 的白皮書中。其設計目標（如下面的兩分鐘視頻中所總結的）值得稱讚：將重複且容易出錯的任務自動化，將最佳實踐進行封裝，保證了從訓練到部署的順利執行，所有這一切都以 scikit-learn 風格進行封裝。

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

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

這些函數被 tf.estimator.Estimator 類調用並返回一個初始化的估算器。通過此估算器，可以調用 .train、.eval和 .predict 函數，而不用關心圖和會話，這兩個組件在基礎的 TensorFlow 設置中比較難用。

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

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

估算器面臨的挑戰

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

因為其良好的默認檢查點和 Tensorboard 集成，估算器在訓練中使用起來很方便。然而，我們認為推斷的介面有點不大直觀。

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

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

「TensorFlow 估算器：在高階機器學習框架下實現間接性和靈活性」，第 4 頁，作者  Cheng  等人

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

TF 估算器方法的約定

.train、.eval、.predict 都會用到 tensorflow 稱為 input_fn 的函數。調用此函數會返回一批數據。

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

這意味著對於估算器而言，訓練循環是在內部進行的。這樣做很有道理，正如白皮書中所強調的：

因為訓練循環非常普遍，對其的最好實現應該是移除許多重複的用戶代碼。這在理論上很簡單，我們可以避免由此產生的一些錯誤，不讓用戶為此而煩惱。因此，估算器實現並控制了訓練循環。

「TensorFlow 估算器：在高階機器學習框架下實現間接性和靈活性」，第 5 頁，作者  Cheng  等人

這樣的設計可以很好地滿足需要預先對送入估算器的數據進行指定的情況。該使用場景常出現在訓練和評估中。

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

原始的推斷

假設我們想要將訓練過的估算器用於另外一個任務，同樣是使用 Python。我們通常希望在一個工作流程中組合使用多個模型，例如使用語言模型作為自動語音轉錄或光學字元識別中定向搜索的補充。

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

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

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

估算器的 .predict 方法已經被封裝，所以調用 FlowerClassifier.predict() 會返回一個經過訓練的估算器的預測值。

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

緩存估算器來推斷

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

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

乍看起來不大直觀，我們通過一個例子仔細研究一下到底發生了什麼：

[主線程]： 用戶調用 .predict 方法

[主線程]： 將一系列新的數據被添加到 input_queue

[輔助線程]：數據生成器將從 input_queue 中生成一個輸入實例

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

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

[主線程]: 調用封裝好的模型，返回 output_queue 中的最新項

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

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

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

代碼

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

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

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

想知道更多深度學習的技巧，訂閱  Element AI Lab Blog。

雷鋒網雷鋒網

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！