TensorFlow + MKL 内存泄漏及解决办法

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，但是内存确实已不再增长。

结论

1. TensorFlow c++ + MKL 确实是会导致内存上涨，不是泄漏，但是内存确实是会慢慢上涨
2. 内存上涨与调用次数有关
3. TensorFlow CPU本身并不会造成内存增长