1. 從 TF1 到 TF2, 線上內存爆炸了

最近我們團隊使用的框架從 TF1 升級到了 TF 2。升級之後，線上的 Tensorflow Serving 發生了爆內存的現象。具體現象如下圖所示：16G 的內存不到半個小時全部耗盡，內存耗盡之後服務掛掉，然後服務管理平臺重新拉起服務；不到半個小時，16G 內存又耗盡，服務掛掉又拉起；這個過程反覆進行。TensorFlow Serving 進程因 Out-of-Memory 多次重啓。

之前，團隊使用 TF1 tf.feature_column + tf.estimator 的組合編寫訓練代碼，並基於 TensorFlow Serving 搭建模型推理服務。團隊採用這一技術方案，有效地滿足了諸如企鵝電競、王者營地、掌上英雄聯盟等業務的需求。

近日，團隊將 Tensorflow 升級到 2 版本。考慮到在 TF2 版本中官方更推薦使用 tf.keras API。於是編寫訓練代碼時，也將 tf.estimator 替換成了 tf.keras。訓練完成後，導出 SavedModel 格式的模型文件，使用 TensorFlow Serving 部署上線。

2. 空間換時間的遺禍

對於 TensorFlow Serving 進程內存佔用暴漲這一問題，我們從排查 TensorFlow 內存泄漏的角度做了些許嘗試，但都無果而終。幸運的是，我在 GitHub 上查閱相關的 issues 時，無意間發現了一個最近的 Pull requests [3]。裏面提到，如果模型有 n 個輸入，那麼，有可能產生 n! 種可能的 key，從而導致佔用大量內存。

帶着這個問題，我閱讀了 direct_session.cc 中的 [GetOrCreateExecutors] 函數[1]和 predict_util.cc 中的 [PreProcessPrediction] 函數[2]。

原來，Session 在初次執行時，會對模型輸入、輸出的 key_name 做排序並拼接，得到 sorted_key，然後創建相應的 Executor，並將 <sorted_key: executor> 這一映射保存到 <key: executor> 字典中。後續 Session 在執行時，會先判斷是否已經存在與模型輸入、輸出對應的 Executor，若存在，則重複使用這個 Executor；若不存在，則創建一個新的 Executor。

但是，TensorFlow Serving 是通過遍歷一個 Map 來從 PredictRequest 中抽取模型輸入、輸出的 key_name，因此，每次遍歷的順序都不盡相同。理論上，假如模型有 n 個輸入，那麼有 n! 種可能的遍歷結果。

而 TensorFlow 恰巧在查找 Executor 上做了以空間換時間的“優化”。它並不會每次都對輸入、輸出的 key_name 做排序、拼接，而是先直接拼接，得到 unsorted_key，並查找 <key: executor> 這個字典，若未找到，再對輸入、輸出做排序、拼接， 然後查找。同時，將 <unsorted_key: executor> 這一映射保存到 <key: executor> 字典中，以提高後續查找 Executor 時的“命中率”，避免排序。

那麼爲什麼在 TF1 的時候，沒有這個問題呢？因爲 不同於 TF1 的 tf.feature_column + tf.estimator 的組合，使用 TF2 的 tf.feature_column + tf.keras 的組合導出的模型有多個輸入。 下面的截圖分別展示了兩種組合的（部分）輸入。

左邊對應 TF1 的 tf.feature_column + tf.estimator，模型僅有一個輸入（protobuf序列化後的tf.Example）；右邊對應 TF2 的 tf.feature_column + tf.keras，模型有多個輸入（每個輸入對應一個特徵）。 考慮到我們 模型 的輸入特徵數量接近 400 ，400! 是一個天文數字，足以塞爆內存 。

基於上面調查的結果，我按照上述 [Pull requests][3] 修改 [PreProcessPrediction] 函數[2]的源碼，每次獲得輸入、輸出的 key_name 後都對其進行排序。 重新編譯 tensorflow_model_server 程序，再次對服務做壓力測試，下面的截圖給出了壓測過程中內存使用率的變化曲線。

可以看到，除了壓測剛開始的一段時間，其餘時間 TensorFlow Serving 內存佔用保持平穩。

3. 不夠恰當的優化

空間換時間是很好的優化方法。但 TF 用空間換時間來優化輸入（特徵） 的排序，卻很無厘頭。如果輸入（特徵）少，比如只有 10，省掉排序能省下多少時間呢？如果輸入（特徵）稍微多一點，比如20，50 和 100，省掉排序能省下的時間多了一點，但 20！= 2432902008176640000，50！和 100 ! 都是天文數字，很容易塞爆內存。

空間換時間本身是一個很好的方法，但使用的時候需要特別注意，會不會使用很多內存從而爆內存。如果使用很多的內存，則需要採取一定的限制措施。比如在 TF 輸入（特徵） 的排序場景。如果能用一個容量上限的  map，並且超過這個 map 容量上限就將最近最少使用的  key 淘汰 ，那麼既不會爆內存，時間也能得到優化。

4. 結論

首先膜拜下 Pull requests [3] 的作者。對 TF 足夠深入的理解，纔能有這種發現。

應該不少團隊準備或者正在將 TF1 升級到 TF2，而使用的模型有多個輸入（特徵），就很有可能碰到內存爆漲的問題。畢竟一個模型超過 20 個輸入（特徵）是非常正常的事情，20! = 2432902008176640000 已經是一個足以塞爆內存的天文數字。

不要等 TF 官方合併Pull requests [3] ，我們自己把這個 Pull requests [3] 的改動引入之後重新編譯 TF Serving，就能解決問題。或者自己實現一個 “用一個固定容量的  map，並且超過這個 map 容量就將最近最少使用的  key 淘汰 ”，然後給 TF 再提一個 PR。

5.    鏈接

[1] https://github.com/tensorflow/tensorflow/blob/463c3055ecd3bba92d7e1da3ebe48e7e8394a0c1/tensorflow/core/common_runtime/direct_session.cc#L1455

[2]

https://github.com/tensorflow/serving/blob/99dfd14c11cfd500ff3beda0c8f1b99094e9d88d/tensorflow_serving/servables/tensorflow/predict_util.cc#L89

[3]

https://github.com/tensorflow/serving/pull/1638