TensorFlow2内存泄漏问题优化
一、前言
TensorFlow2是相对于TensorFlow1发生了较大改变，从TensorFlow1的静态图计算改成TensorFlow2的动态图计算。TensorFlow1因为代码结构定义复杂，需要先构建图再执行图，计算效率比TensorFlow2要高，缺点也比较明显代码写起来很麻烦。谷歌估计也是受了Pytorch框架的刺激，被吐槽tf框架难用，在TensorFlow2的时候进行了大改进将原来的静态图改成动态图，写代码变得简单了但同时也带来了内存泄漏问题，尤其是在线上部署时随着推理的不断增多推理程序所在的容器也是不断增加，当达到一定程度时就会被系统强制kill掉。内存泄漏bug一直持续到最新版的tf也没有解决，我下面提出的一些方法也不是从框架本身的内部bug解决，更多的是通过一些方式绕过。
二、内存泄漏原因探究
线上部署的服务从以前的TensorFlow1全面升级成TensorFlow2，在进行上线压测时发现随着推理样本增加容器占用的内存问题也是线性增加，线上推理项目是三元组抽取，采用6层roberta进行cpu推理，模型占用的也是内存空间，进行一晚上压测容器内存占用也从原来的1g迅猛扩大到10g+，如果是这样线性增长下去迟早会被系统kill掉，也不敢放在线上服务使用。
内存增长
随后进行原因探查，首先找到tensorflow issue里面提到的类似问题：
https://github.com/tensorflow/tensorflow/issues/54086
但这个issue并未给出解决方案，官方有回复但只是说收到反馈之类的。
我上网查了一些TensorFlow2相关的问题，也有不少人或者公司遇到过内存泄漏问题。
网上言论
我尝试定位原因，内存增长主要发生在模型侧，随着每次推理模型那块的内存都会增长一些，更加确定是模型预测那块导致内存增长。无论是tf1还是tf2都是读取save model格式PB模型来推理，加载方式略有不同。
TensorFlow1模型加载方式
def ner_model(dict_path,pb_path):
"""
命名实体识别模型构建以及加载权重.
:param
dict_path:str
pb_path:str
:return sess,tokenizer 返回模型会话和bert预处理分词器
"""
tokenizer = Tokenizer(dict_path, do_lower_case=True)
sess = tf.Session(graph=tf.Graph())
tf.saved_model.loader.load(sess,[tf.saved_model.tag_constants.SERVING], pb_path)
return sess,tokenizer
TensorFlow1模型推理方式
#事先定义模型图层输入输出名
x1 = self.sess.graph.get_tensor_by_name('Input-Token:0')
x2 = self.sess.graph.get_tensor_by_name('Input-Segment:0')
y = self.sess.graph.get_tensor_by_name('global_pointer/truediv:0')
scores_pred = self.sess.run(y,
feed_dict={x1: batch[0], x2: batch[1]})
TensorFlow2模型加载方式
def ner_model(dict_path,pb_path):
"""
命名实体识别模型构建以及加载权重.
:param
dict_path:str
pb_path:str
:return sess,tokenizer 返回模型会话和bert预处理分词器
"""
tokenizer = Tokenizer(dict_path, do_lower_case=True)
sess = tf.saved_model.load(pb_path)
return sess,tokenizer
TensorFlow2模型推理方式
#无需实现定义输入输出图层
inp0 = tf.constant(batch[0],dtype=tf.float32, name='Input-Token')
inp1 = tf.constant(batch[1],dtype=tf.float32, name='Input-Segment')
scores_pred = self.sess([inp0,inp1]).numpy()
**静态图（**需要先构建再运行）：
- 优势是在运行前可以对图结构进行优化，比如常数折叠、算子融合等，可以获得更快的前向运算速度。
- 缺点也很明显，就是只有在计算图运行起来之后，才能看到变量的值，像TensorFlow1.x中的session.run那样。
**动态图（**一边运行一边构建）：
- 优势是可以在搭建网络的时候看见变量的值，便于检查。
- 缺点是前向运算不好优化，因为根本不知道下一步运算要算什么。
所以TF2中在推理时是不需要定义输入图层，会通过输入值在计算构建计算图，所以TF2内存泄漏最主要的原因是：
不同的输入会导致模型推理时会一直构建图节点，经过大量样本输入时会让计算图变得非常大，所以内存一直在线性增长也就是因为计算图不断在缓存。
三、内存泄漏解决办法
根据上面的探查其实我们已经找到原因，就开始着手解决问题。
以下是我提出的一些解决方法：
1、使用TF2的兼容模式读取模型并且关闭动态图（本人线上解决办法）
现在知道内存叠加也主要因为动态图缘故，所以我们在推理的时候也需要关闭动态图，并且通过TF2的兼容模式可以加载TF1训练的模型在线推理。
#TF2兼容模式
import tensorflow.compat.v1 as tf
#关闭动态图
tf.disable_v2_behavior()
经过此方法线上docker占用内存进入一个稳定的状态
线上内存检测
在尝试上述方法前我也使用过一些操作但是并未阻止内存增加，仅仅只能延缓内存增长速度，但也不是完美的解决办法。
#清空计算图
tf.reset_default_graph()
#销毁当前计算图
tf.keras.backend.clear_session()
#显示删除变量用gc模块回收
gc.collect()
2、定时剔除不常用的节点
3、定时重启容器或者检测内存重启容器
线上有些定时会重启容器避免内存无限制增长或者检测内存重启容器。
四、总结
只能说一句，TF太坑啦！！！
