调用Keras模型.fit()后出现内存泄漏。

Question

我编写了自定义的Keras层，称为Material，并尝试在我的模型中使用它：

import tensorflow as tf
from tensorflow import keras
import numpy as np

class Material(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.table = tf.Variable(
            initial_value=np.ones(shape=(6, 8, 8), dtype="float32"),
            trainable=True,
        )
        self.k = tf.Variable(initial_value=8., trainable=True, dtype='float32')
    
    def call(self, inp):
        material_white = tf.reduce_sum(inp[..., 0] * self.table, axis=(-1, -2, -3))
        material_black = tf.reduce_sum(inp[..., 1] * self.table, axis=(-1, -2, -3))
        material_white = tf.maximum(material_white, .01)
        material_black = tf.maximum(material_black, .01)
        return tf.math.log(material_white / material_black) * self.k

def get_model() -> keras.Model:
    inp = keras.Input((6, 8, 8, 2),)
    material = Material()
    out = keras.activations.sigmoid(material(inp))
    return keras.Model(inputs=inp, outputs=out)

model = get_model()
model.compile(optimizer=keras.optimizers.Adam(learning_rate=.01), loss='mean_squared_error')

board_tables = np.random.sample((100000, 6, 8, 8, 2),)
outcome = np.random.sample((100000,),)

while 1:
    model.fit(board_tables, outcome, batch_size=200, epochs=2)

由于某种原因，在每次fit迭代（我指的是函数调用，而不是每个时期）后，它似乎会使用越来越多的RAM。在大约5GB时，内存使用增长减缓但仍然持续存在。这个问题在CPU和GPU上都存在。有人能解释一下发生了什么吗？我的自定义层有什么问题吗？

感谢任何建议。

- Усердный бобёр

Answer 1

我使用了从自定义生成器创建的tf.data.Dataset来将数据传递到我的模型中。出于某种原因，内存泄漏现在已经消失了。但我仍然不明白普通的np数组有什么问题。

- Усердный бобёр

---

我也遇到了同样的问题。你能提供你的解决方法吗？我可以确认这个问题确实存在，而且clear_session()一点用都没有。 - Alessio Mora

@AlessioMora 我的代码太长且难以阅读，我觉得在这里贴出来不可行。只需编写一个生成器，它会产生(x, y)，其中x和y是模型输入和期望输出的批次。然后执行dataset = tf.data.Dataset.from_generator(your_generator, output_types=(tf.float32, tf.float32))（当然要使用正确的数据类型）。接下来执行model.fit(dataset, epochs=...)。你还可以在fit方法中提供steps_per_epoch参数，以便在特定批次数之后停止训练，而不是等到生成器用尽。 - Усердный бобёр

我不会创建新的数据集。我只是制作了一个数据集生成器，可以无限地产生我想要的任何数据。 - Усердный бобёр

Answer 2

如果您只增加epochs的数量而不是将model.fit放入while循环中，是否会发生同样的情况？不会吗？

由于您在无限循环（while 1）中重复训练模型，因此随着时间的推移，内存使用量会增加。

尝试在每个model.fit迭代之后清除会话：

keras.backend.clear_session()

- mrk

---

这没有意义，这就是所有Keras层的实现方式，权重变量只存储一次（在类内部）。 - Dr. Snoopy

你试过了吗？当我在循环中多次调用model.fit函数时，一段时间后会出现问题，就像问题中所报告的那样。但是，当我清除会话时，这个问题就消失了。这可能不是变量的问题，我承认这一点（已编辑我的答案）。 - mrk

keras.backend.clear_session() 不会改变任何东西。如果我只是增加 epoch 的数量，它也不会发生变化。但是我需要在不同的数据上调用 model.fit，所以我实际上无法避免使用循环。 - Усердный бобёр