TensorFlow + MKL 內存泄漏及解決辦法
背景
最近在做深度學習在線推理相關的項目，底層使用TF來作爲推斷框架，我們知道，TF底層是用Eigen完成CPU的計算的，而Eigen的速度只能說一般般，與成熟的計算庫在性能上的差距還是比較明顯的，比如TF支持的intel MKL。
之於如何使用TF+MKL，可以參考此文，其實就是在編譯時加上--config=mkl選項，然後在生成可執行文件時鏈接libmkldnn.so, libmklml_intel.so, libimop5.so，就可以使用TF+MKL來進行推斷了。
聽上去很美好，實際用起來也很絲滑，但是當上到服務器上的時候，就出問題了！
將算法封裝成Java服務，JNI調用算法接口進行推斷。使用壓測工具模擬10個併發的壓測，半個小時後該進程佔用服務器內存過多導致進程被kill。
原因分析
是什麼原因呢？
首先從內存暴漲想到內存泄漏！前面說半個小時內存暴漲並掛掉，實際上26分鐘就掛掉了，壓測工具顯示共進行了86000次調用，16G的內存從剛開始調用的15%漲到100%，粗略計算一下，平均每次調用上漲0.1619Mb。如果有泄漏，那應該是非常明顯的了，畢竟意味着165Kb的數據被申請並且沒有被釋放。
查！
內存泄漏
這裏我們有必要說一下C++代碼內存泄漏的檢查方法，首先是模擬調用過程，100000次調用，10個線程組成線程池來調用算法接口，使用top命令，觀察可執行程序所在進程的物理內存數量，如果物理內存一直上漲，那麼存在內存泄漏。
使用以上方法，測得可執行程序在運行期間，物理內存大概每10分鐘漲10M，這個數字不多，是因爲我本機運行慢，說明兩點，一是內存上漲明確和多次調用有關，二是服務器單次運行快，所以內存漲的也快。
具體是哪裏出了問題？爲什麼能確保就是TF甚至能定位是MKL的問題呢？
另外一個粗暴的辦法是註釋，註釋掉session->run()及後面的代碼，內存不上漲；然後只打開session->run()，有內存上漲，基本確定和session->run有關。
接下來是valgrind工具，使用方法參考此文，當我們對可執行文件使用valgrind工具進行檢測時，大喫一驚：
MKL的內存泄漏
爲什麼MKL就有內存泄漏呢？
這個issue給出了一個可能的解釋，摘出一小段，有時間建議閱讀全文：
Just for trouble-shooting purpose, please experiment by hard-coding
do_not_cache = true; (comment out existing logic).
My initial feeling is that more and more MKL primitive caching due to new op input dimensions (thus new keys), so that the memory usage is going higher and higher (that is, the feeling of leakage).
Of cause, this is just a guess. To confirm it, do the hard-coding experiment on all do_not_cache
occurrences (in files of mkl*.cc in folder tensorflow/core/kernels).
大意就是說，並沒有發生真正的泄漏(即指向某段內存的指針失效或者指向別處)，而是每次調用都新建了一塊內存，且這個指針被保存起來，進程終止時釋放，實際表現也和這個類似，雖然沒有發生真正的泄漏，但是內存確實會因此一直上漲進而導致崩潰。
如何解決？
首先我們參考MKL自己的說法，並沒有什麼用。。。
只能試一下去掉MKL，如果內存還是上漲，那就是TF本身的問題了，但是我不相信會有這樣的事情發生。這裏順帶說一下，在編譯tf的c++的庫的時候，可能會在編譯某個op的時候報錯，不要慌張，兩步解決，一是增加機器的內存，8G最好，二是在bazel build命令後面加上--config=opt，實測有效。
在去掉MKL之後，同樣的驗證方式，內存不再上漲，問題解決：
這裏實際還是有一些possibly lost和still reachable，但是內存確實已不再增長。
結論
- TensorFlow c++ + MKL 確實是會導致內存上漲，不是泄漏，但是內存確實是會慢慢上漲
- 內存上漲與調用次數有關
- TensorFlow CPU本身並不會造成內存增長