如何使用 K.clear_session() 修复 Keras 中的内存泄漏问题？答案

【问题标题】：How to fix memory leak issue in Keras using K.clear_session()?如何使用 K.clear_session() 修复 Keras 中的内存泄漏问题？
【发布时间】：2020-11-25 05:17:58
【问题描述】：

我有一个网络，我正在通过提供批量数据来训练它，我正在使用model.train_on_batch() 来执行此操作。如果我只运行这个训练部分，我会看到我的网络在 40 多个 epoch（到目前为止）的 RAM 利用率为 3% 时训练得很好，每个 epoch 大约有 2000 次迭代。当我尝试在每个 epoch 之后进行验证（也分批发生）时，内存泄漏非常严重，导致 90% 的 RAM 利用率和我的代码被杀死。所以我在过去几天尝试了几件事，似乎model.predict() 在循环中导致内存泄漏open issue at Tensorflow GitHub。我试过predict_on_batch()，同样的行为。 model(inputs, training=False) 似乎减缓了内存泄漏，而不是从 3% - 7% - 13% - 40% - 80% - 90% （间隔 60 秒）突然跳跃，它增加了每分钟 1%。但在某个时候，它也达到了 90%。我唯一可以尝试这个 github 线程的是使用 K.clear_session()。

我尝试阅读 K.clear_session() 的文档和一些 SO 帖子，都建议在创建多个模型时使用它，我没有这样做。所以我的问题是，如果我有一个模型在循环中被训练和评估，我应该在哪里使用K.clear_session()，在每个时期之后并在每个时期之前重新加载保存的模型？那是正确的吗？

除此之外，我还得到了拓扑排序错误another open issue，所以我想知道是不是因为我正在循环训练，因为我的代码没有循环，这也导致了内存泄漏，@987654330 @ 会有所帮助吗？

我的代码结构的最小示例：

from tensorflow.keras.models import Model
K = tf.keras.backend

def myModel():
    **some architecture**

ip = Input(shape=(h, w, 3))
op = myModel(ip)
model = Model(ip, op)
model.compile(optimizer=Adam(lr=1e-6), loss=custom_mean_squared_error)

for e in range(numEpochs):
    for batch in range(0, num_train_batches):
        x = readImages()
        y = readLabels()
        loss = model.train_on_batch(x, y)
        

    for batch in range(0, num_val_batches):
        x = readImages()
        y = model.predict(x)
        val_loss = K.get_value(custom_mean_squared_error(x,y))
        # save predictions

# plot training vs validation loss

Tensorflow-gpu-1.14，python3.6。如果我也做错了什么，将不胜感激。

【问题讨论】：

标签： python tensorflow keras

【解决方案1】：

这似乎对我有用，但会使过程变慢：

from tensorflow.keras.models import Model
K = tf.keras.backend

def myModel():
    **some architecture**

ip = Input(shape=(h, w, 3))
op = myModel(ip)
model = Model(ip, op)
model.compile(optimizer=Adam(lr=1e-6), loss=custom_mse)

for e in range(numEpochs):
    for batch in range(0, num_train_batches):
        x = readImages()
        y = readLabels()
        # define appropriate flags for first loop
        model = tf.keras.models.load_model(model_path,custom_objects={ 'custom_mse': custom_mse } )
        loss = model.train_on_batch(x, y)

    model.save(model_path)

    for batch in range(0, num_val_batches):
        x = readImages()
        model = tf.keras.models.load_model(model_path,custom_objects={ 'custom_mse': custom_mse } )
        y = model.predict(x)
        K.clear_session()
        val_loss = K.get_value(custom_mean_squared_error(x,y))
        # save predictions

# plot training vs validation loss

【讨论】：