用TensorFlow實現ML模型並調優：每秒可做3億次預測

機器之心發表於2021-09-23
模型

TensorFlow 是目前使用最廣泛的機器學習框架之一，它加快了研究速度，並減少了新模型的生產時間。在一篇論文中，來自原生程式化 DSP 公司 Zemanta 的資料科學總監 Davorin Kopič和工程師 Jan Hartman 展示了將線上廣告生態系統中的大規模機器學習模型轉換為 TensorFlow 框架的過程，並將在 TensorFlow 框架中實現的機器學習模型擴充套件到每秒超過 3 億次預測。因此，該研究的主要內容是在 TF 中實現模型並使用各種最佳化技術以低延遲有效地為其提供服務。


用TensorFlow實現ML模型並調優：每秒可做3億次預測

論文地址：https://arxiv.org/abs/2109.09541


該研究使用的案例是線上廣告的點選預測。在 RTB （實時競價）中，多個 DSP（競標者）透過在網頁的載入過程中實時競標來競爭線上廣告空間。廣告空間是按廣告印象出售的，這使得以市場價值出售虛擬廣告空間成為可能。透過使用機器學習，RTB 還使廣告商能夠最大化其 KPI，例如點選率 (CTR)。估算廣告的點選率是 RTB 的核心問題之一，擁有一個好的點選預測模型非常重要。

在 Golang 中實現的基於自定義邏輯迴歸和分解機 (FM) 的模型，其表達能力受限，並且需要手動實現所有學習程式，這些都會減慢實驗速度並限制模型的預測效能。因此，研究者決定採用 TensorFlow 框架，並用表達能力更強的模型替換現有模型。


面臨的挑戰

基於廣告競標這一具體用例，該研究遇到了一些挑戰，共分為實現、服務和最佳化三個方面。

一方面，為每臺機器配備一個或多個頂級 GPU 的成本會高得令人望而卻步；另一方面，只有一小群 GPU 的機器將迫使研究過渡到基於服務的架構。鑑於這兩種選擇都不是特別可取，而且該研究的模型與其他深度學習領域（例如計算機視覺或自然語言處理）的 SOTA 模型相比相對較小，因此該研究不在生產中使用 GPU 進行推斷。並且由於該研究的模型使用稀疏權重，其用例也不適合 GPU 工作負載。


實現

為了給 TF 模型實現有效的訓練迴圈，該研究實現和測試了多種方法。高吞吐量線上訓練和在 TF 中服務的案例研究很少，而且文件往往不夠具體，這迫使研究者必須通讀原始碼和基準原型才能發現實現過程中的陷阱。

TF 提供了一個龐大的生態系統和大量具有 SOTA 演算法實現的庫。選擇功能豐富的已有實現非常容易，但是研究者發現這些實現大多未經最佳化，因此他們決定自己實現演算法。TF 具有不同抽象級別的 API，不過，一些 API 雖然易於使用，但通常是效率低下的最底層操作（例如 Estimator API）。研究者最終選擇了 Keras3，因為它是一個圍繞底層 TF 操作的瘦包裝器（thin wrapper），並具備高水平的效能，且易於理解。由於 TF 是一個功能和資源都非常豐富的庫，因此該研究還必須考慮要在其中實現多少機器學習 pipeline。研究者選擇暫時擱置特徵轉換和互動，僅實現學習演算法——儘管它們是最小的可替換部分，但卻具備最大的改進潛力。

由於 Golang TF 包裝器僅支援預測，因此必須在 Python 中實現訓練迴圈。指令碼透過將其標準輸入作為子程式實現與 Golang 資料 pipeline 的連線。資料以高效的二進位制格式傳送而無需解析，與 CSV 格式相比，該方法的速度提高了 25%。然後在後臺執行緒中讀取資料，以防止模型在等待資料時空閒。基於此，該研究實現了在整個訓練 pipeline 中保持高吞吐量。事實證明，高效的輸入和輸出也是低延遲預測的關鍵，該研究透過將所有輸入特徵連線到單個張量（tensor）中，顯著減少了在序列化和複製輸入資料上花費的時間。

用TensorFlow實現ML模型並調優：每秒可做3億次預測


服務

研究者發現，由於計算密集型神經網路的存在，在使用 Golang TF 裝飾器的情況下，DeepFM 模型的 CPU 使用率要高得多。儘管帶來了指標的顯著提升，但將這種方法擴充套件到 100% 的流量會帶來大量的硬體成本。由於當前全球面臨晶片短缺的問題，這意味代價是困難和昂貴的。


顯然，降低計算成本是很有必要的。然而縮小神經網路模型規模的同時，也會降低模型的預測效能。在深入研究了 TF 之後，研究者發現如果在計算批處理時增加示例的數量，計算效率會大大提升。這種低線性增長是由於 TF 程式碼被高度向量化了，TF 也會產生每次計算呼叫的開銷，然後將其分批攤銷。考慮到這一點，如果希望減少計算呼叫的數量，就需要將許多請求連線到一個計算中。

研究者構建了全部包含在一個執行的 bidder 例項的自動批處理系統，以避免網路呼叫。由於每個例項每秒鐘接收數千個傳入請求，因此可以保證連線來自眾多請求的計算，建立更大的批次。研究者透過一些批處理執行緒實現了這一點，這些執行緒接收來自傳入請求的資料，建立批處理並在批處理完成後初始化計算。計算過程中，每隔幾毫秒即初始化一次以避免超時，因為批處理可能在這個時間視窗中沒有填充。這種實現是高度最佳化的，將計算呼叫的數量減少到五分之一，同時將 TF 計算的 CPU 佔用量減半。

雖然在批處理器執行緒沒有獲得 CPU 時間的極少數情況下，會發生請求超時，但只有不足 0.01% 的請求會出現這種情況。研究者觀察到平均延遲略有增加（平均約 5 毫秒），流量高峰時可能會更多一些。因此他們實施了 SLA（服務等級協議）和適當的監控手段，以確保延遲的穩定性。鑑於沒有大幅增加超時的百分比，這些方法仍是非常有效的，也是這一 TF 服務機制的核心。

用TensorFlow實現ML模型並調優：每秒可做3億次預測
本文作者之一 Davorin Kopič

最佳化

研究者在 TF 中實施的模型最初比定製的 FMs 慢得多，為了尋找加速空間，研究者大量使用內建的 TF 分析器來尋找執行時間最長的操作，並儘可能進行了改進。最常見的是各種冗餘的 reshape 或轉換運算。其中一個更有趣的發現是 Adam 最佳化器比 Adagrad 慢得多 (大約 50%)，儘管二者運算數量上的差異很小。分析器顯示，對稀疏權值進行梯度更新需要大量計算時間。這是因為模型的權重是稀疏的 (特徵大部分是分類的，因此非常稀疏) ，而最佳化器沒有考慮到這個事實。

由於用 Adagrad 替換 Adam 意味著深度模型效能的顯著降低，研究者也尋找了其他解決方案：切換到 Lazy Adam 的最佳化器被證明是非常有效的，因為它可以非常有效地處理稀疏權重問題。結果顯示，其整體加快了超過 40% 的訓練速度，與 Adagrad 相接近。

由於使用了自適應最佳化器（比如 Adam），這也需要儲存權重矩和方差，每個引數將儲存三個值，將儲存的模型大小增加了三倍。然而，這些值實際上並不用於預測，只用於訓練。研究者利用這一點構建了最佳化過程，去掉了這些值的模型，減少了 66% 的資料量，並降低了記憶體使用量和成本。

相關文章