由于循环调用tensorflow.keras.predict导致的内存泄漏问题

问题描述

之前在进行强化学习模型训练的过程中无意间发现了这个问题，于是有了这个博客来记录一下。

首先就是出现内存泄漏的一段代码：

import tensorflow as tf
import numpy as np

# 首先简单搭个模型
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(128,)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(128, activation='relu')
])

while True:  # 不断进行循环测试溢出
    inputX = np.random.randn(128 * 10).reshape((10, 128))
    model.predict(inputX)

初始内存如下：

随着循环的进行，pycharm占用的内存随之增加。

但是在上面的代码中并没有需要增加内存占用的部分。

可以看出model.predict的反复执行导致了系统内存泄漏。

解决方法

我查阅了许多的相关资料，有的说将模型输入X转为tensor(张量)在抛进模型的；有的说调用。tf.keras.backen.clear_session的。但经过我的尝试都是无济于事的。

最终经过我的尝试，发现将predict换成predict_on_batch可以有效地解决这个问题，同时在执行速度上来说predict_on_batch也是比predict要快的。