tf.keras model.predict 导致内存泄漏

Posted 2023-02-23

【中文标题】tf.keras model.predict 导致内存泄漏

【英文标题】：tf.keras model.predict results in memory leak

【发布时间】：2021-01-19 18:34:26

【问题描述】：

在谷歌 colab 上工作。使用 tf.keras 和 tensorflow 版本 2.3.0 我快疯了，因为我无法使用我训练过的模型来使用model.predict 运行预测，因为它耗尽了 CPU RAM。我已经能够通过一个非常简单的示例重现该问题。

import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Input,Conv2D, Activation

matrixSide = 512 #define a big enough matrix to give memory issues

inputL = Input([matrixSide,matrixSide,12]) #create a toy model
l1 = Conv2D(32,3,activation='relu',padding='same') (inputL) #120
l1 = Conv2D(64,1,activation='relu',padding='same')(l1)
l1 = Conv2D(64,3,activation='relu',padding='same')(l1)
l1 = Conv2D(1,1,padding='same')(l1)
l1 = Activation('linear')(l1)
model = Model(inputs= inputL,outputs = l1)


#run predictions
inImm = np.zeros((64,matrixSide,matrixSide,12))
for i in range (60):
  print(i)
  outImm = model.predict(inImm)
# K.clear_session() #somebody suggested it...

基本上，在 GPU 上工作时，它在前 4 次迭代中使用 3.0 GB 的 CPU RAM，然后它上升到 7，然后到 10，然后它崩溃了，因为它耗尽了所有可用的 RAM！ 在 CPU 上运行时，它会持续进行更多迭代，有时它甚至会将其使用的 RAM 量从 9 GB 减少到 3 GB，但最终在 20 次左右的迭代后它仍然崩溃。

上一个示例 (Keras predict loop memory leak using tf.data.Dataset but not with a numpy array) 在使用 tf.data 时也有类似的问题，但在使用 numpy 时没有。有人在 github 问题上建议 tensorflow 1.14 在每个循环中执行 K.clear_session ......但这没有帮助！

知道如何解决这个问题吗？

【问题讨论】：

如果这是 TF 1.x - 添加正确打开和关闭会话的命令或使用 with 上下文：***.com/questions/53885356/…

不，我使用的是 TF 2.x

我也遇到了这个bug

我也遇到过这个问题。我能做的最好的事情：在此处收集解决方法列表：Keras memory leak

【参考方案1】：

我正在使用基于keras docs的简单解决方案

对于适合一个批次的少量输入，建议直接使用 call() 以加快执行速度，例如 model(x) 或 model(x, training=False)

for filename in image_filenames:
  # read of data
  input = load_image(filename)

  # prediction
  output = model(input) # executes __call__() or call()

使用__call__() 或model(input) 避免了predict 方法内部的内存泄漏，该方法每次执行都会创建一个数据生成器，并且不会释放内存。

【参考方案2】：

我通过使用K.clear_session() 解决了这个问题。 首先，您需要先定义一个会话，然后才能清除它。 这样做的目的在here 和here 中都有说明。

config= tf.ConfigProto(log_device_placement=True) 
config.gpu_options.allow_growth = True
session = tf.Session(config=config)
K.set_session(session)

起初，在循环中使用K.clear_session() 会导致第一次预测后出现错误。在我看来，tf 失去了与 model 的连接。 出于这个原因，我在循环的每次运行中创建一个新模型。这会对前多次运行的代码速度产生负面影响，但会阻止 RAM 存储的累积。

以下代码包含建议的改进：

import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Input,Conv2D, Activation

matrixSide = 512 #define a big enough matrix to give memory issues

config = tf.ConfigProto(log_device_placement=True)
config.gpu_options.allow_growth = True
session = tf.Session(config=config)
K.set_session(session)

def create_model(matrixSide_v):
    inputL = Input([matrixSide_v,matrixSide_v,12]) #create a toy model
    l1 = Conv2D(32,3,activation='relu',padding='same') (inputL) #120
    l1 = Conv2D(64,1,activation='relu',padding='same')(l1)
    l1 = Conv2D(64,3,activation='relu',padding='same')(l1)
    l1 = Conv2D(1,1,padding='same')(l1)
    l1 = Activation('linear')(l1)
    c_model = Model(inputs= inputL,outputs = l1)
    return c_model

#run predictions
inImm = np.zeros((64,matrixSide,matrixSide,12))
for i in range(64):
    print(i)
    model = create_model(matrixSide)
    outImm = model.predict(inImm)
    K.clear_session()

【讨论】：

在我的用例中，我使用经过训练的模型来预测大量样本。为每个预测加载模型并清除会话是没有意义的，因为它既慢又慢，并且加载模型也有记录的内存泄漏问题。任何建议如何将您的解决方案应用于我的用例？

我同意，每次迭代重新加载模型会使事情变慢 10000 倍，并不是真正的实用解决方案。

【参考方案3】：

这是我将其作为错误发布到 Tensorflow 后的理解。

将代码更改为;

in_imm = np.zeros((64,matrix_side,matrix_side,12))
for i in range (60):
  print(i)
  tensor = tf.convert_to_tensor(in_imm, dtype=tf.float32)
  out_imm = model.predict(tensor)

在带有 numpy 输入的 for 循环中使用 tf.keras.Model.predict 每次迭代都会创建一个新图，因为 numpy 数组是使用不同的签名创建的。将 numpy 数组转换为张量保持相同的签名并避免创建新图。

【讨论】：

谢谢，我试试这个！顺便说一句，这些“图表”实际上是什么以及它们的行为方式是否有明确的解释？如果是图形问题，为什么 K.clear_sessions() 不起作用（但 gc.collect() 起作用）？

【参考方案4】：

我找到了解决内存泄漏的方法。虽然 K.clear_session() 在我的情况下没有做任何事情，但在每次使用 _ = gc.collect() 调用之后添加垃圾收集实际上可以解决问题！ 现在实际使用的内存是恒定的，我可以运行任意数量的预测。