调用 Keras model.fit() 后出现内存泄漏

Question

我编写了名为

Material

的自定义 Keras 层，并尝试在我的模型中使用它：

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（我的意思是函数调用，而不是每个时期）。在 5 GB 左右，内存使用量增长放缓，但仍在继续。而且这个问题在CPU和GPU上都存在。谁能解释一下发生了什么事吗？我的自定义图层有问题吗？

感谢您的任何建议。

Answer 1

如果您只是增加

epochs

的数量，是否也会发生同样的情况 而不是将您的

model.fit

放入 while 循环中？没有？

由于您在无限循环中重复训练模型 (

while 1

)，内存使用量会随着时间的推移而增加。

尝试在每次

model.fit

迭代后清除会话：

keras.backend.clear_session()

Answer 2

我使用从自定义生成器创建的

tf.data.Dataset

将数据传递到我的模型中。由于某种原因，内存泄漏现在已经消失了。但我仍然不明白常规 np 数组有什么问题。

Answer 3

你找到解决办法了吗？我也有同样的问题