[续] 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
暂时没有解决。
记录一下,希望能有用,也希望这个链路能早日畅通!