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

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

背景

之前就发现了tf+mkl的内存泄漏问题，当时给出的办法是不用mkl了，使用原版的TensorFlow C++动态库。这样内存确实不怎么泄漏了(实测还是会有一点点上涨)，但是CPU上的性能也就没有了，毕竟已知在我们的机器上，mkl的4个线程单次推理能够比原版快一倍，从600ms降到300ms，相当可观的性能提升了。

思来想去，还是要把这个坑填上。

填坑之旅

1.重燃希望

本来已经不抱希望的，直到发现百度的paddle库中有一个issue说解决了这个问题：[dev] mkl多线程下存在内存泄露问题 #22827，如逢甘露，马上点击阅读，总结一下：

1.mklml本身不会有泄漏
2.锅给到了libiomp5.so！这个libiomp5.so是什么呢？简单来说，如果我们在编译tf时加上mkl选项–config=mkl，那么默认编译出来的libtensorflow_cc.so会依赖libmklml_intel.so，从名字可以知道，这个库就是mklml加上intel的一个底层库，什么底层库呢？就是这个libiomp5.so，类似于gnu默认的libgomp.so，这个库是intel定制开发的OpenMP的库。

我知道上面这段话会有点晕，来个图吧：

左边是默认，右边是推荐的方案。

2.修改tf编译

找到tensorflow/third_party/mkl/mkl.BUILD文件，里面有：

cc_library(
    name = "mkl_libs_linux",
    srcs = [
        "lib/libiomp5.so",
        "lib/libmklml_intel.so",
    ],
    visibility = ["//visibility:public"],
)

我们将里面的库换成目标库：

cc_library(
    name = "mkl_libs_linux",
    srcs = [
        "lib/libmklml_gnu.so",
    ],
    visibility = ["//visibility:public"],
)

然后执行bazel build就可以了。

用图来说明一下：
这是默认的tf库：

这是默认加上mkl的库：

这是修改之后的库：

在做完上述编译和验证后，我们发现：

完全没有用！

3.最后的倔强

这里要说的也是最终的方法，如果确定要在intel的机器上使用tensorflow C/CXX库来进行推断，可以使用intel_tensorflow，下载最新的release，然后直接编译，得到tf库，还是依赖libiomp5和libmklml_intel，但是内存问题就解决了！

重要的事情再说一遍，解决办法：

1.访问intel_tensorflow
2.下载最新release
3.编译tf库

可以看到这个库是2020年8月30号才更新的，而且这个更新有大量的feature，线程池也重新做了设计，不知道是不是重新设计的线程池有效果，总之内存的问题是消失了。

4.残留

虽然内存的问题解决了，但是引入了新的问题：

2020-10-10 10:37:04.850963: W tensorflow/core/framework/op_kernel.cc:1651] OP_REQUIRES failed at mkl_softmax_op.cc:313 : Aborted: Operation received an exception:Status: 2, message: could not create a descriptor for a softmax forward propagation primitive, in file tensorflow/core/kernels/mkl_softmax_op.cc:310

暂时没有解决。

记录一下，希望能有用，也希望这个链路能早日畅通！