在循环中使用 tf.Model 和 tf.Model.fit() 时发生内存泄漏。 clear_session() 没有帮助 [TensorFlow 2.11.0]

Question

对于 TF 版本 == 2.11.0，在循环中使用 tf.Model 和 tf.Model.fit() 时，内存使用量会稳步增加，并最终导致内存不足异常导致内存饱和。 clear_session() 没有帮助。 TF 版本 == 2.9.2 的相同代码具有几乎不变的内存使用，并且按预期工作。

重现代码：

import tensorflow as tf
import time

class MyModel(tf.keras.Model):

  def __init__(self):
    super().__init__()
    self.dense1 = tf.keras.layers.Dense(1000, activation=tf.nn.relu)
    self.dense2 = tf.keras.layers.Dense(10000, activation=tf.nn.softmax)
    self.dense3 = tf.keras.layers.Dense(10000, activation=tf.nn.softmax)
    self.dense4 = tf.keras.layers.Dense(1000, activation=tf.nn.softmax)

  def call(self, inputs):
    x = self.dense1(inputs)
    x = self.dense2(x)
    x = self.dense3(x)
    x = self.dense4(x)
    return x

for r in range(0, 10000):
    model = MyModel()
    ds = tf.data.Dataset.from_tensor_slices((tf.random.uniform((64*4, 1000)), tf.ones((64*4))))
    model.compile(optimizer='sgd', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))

    model.fit(ds.batch(64))
    tf.keras.backend.clear_session()
    time.sleep(3)
    print("round: ", r)

系统信息：

操作系统平台和发行版（例如 Linux Ubuntu 16.04）：Ubuntu 22.04.1 LTS (GNU/Linux 5.16.10 x86_64)

TensorFlow 安装自（源代码或二进制文件）：source

TensorFlow 版本（使用下面的命令）：2.11.0

蟒蛇版本：3.10.6

Answer 1

我面临着与 TF 2.11 完全相同的问题。使用 tf.keras.backend.clear_session() 和 gc.collect() 适用于 TF 2.9.2。