Keras 在每个 epoch 中占用的内存量无限增加

Question

我正在运行一个遗传超参数搜索算法，它很快就会使所有可用内存饱和。

经过几次测试，看起来 keras 所需的内存量在不同时期和训练不同模型时都会增加。随着 minibatch 大小的增加，问题变得更糟，1~5 的 minibatch 大小至少让我有足够的时间看到内存使用量在前几次拟合中非常快速地上升，然后随着时间的推移缓慢但稳定地增加。

我已经检查了keras 预测内存交换无限期增加，Keras：在进行超参数网格搜索时内存不足，以及Keras（TensorFlow，CPU）：循环中的训练序列模型会占用内存，所以我已经在清除 keras 会话并在之后重置 tensorflow 的图形每次迭代。

我还尝试明确删除模型和历史对象并运行 gc.collect() 但无济于事。

我在 CPU 上运行 Keras 2.2.4、tensorflow 1.12.0、Python 3.7.0。我为每个基因运行的代码和我用来测量内存使用的回调：

import tensorflow as tf
import keras as K

class MemoryCallback(K.callbacks.Callback):
    def on_epoch_end(self, epoch, log={}):
        print(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)


def Rateme(self,loss,classnum,patience,epochs,DWIshape,Mapshape,lr,TRAINDATA,TESTDATA,TrueTrain, TrueTest,ModelBuilder,maxthreads):

K.backend.set_session(K.backend.tf.Session(config=K.backend.tf.ConfigProto(intra_op_parallelism_threads=maxthreads, inter_op_parallelism_threads=maxthreads)))

#Early Stopping
STOP=K.callbacks.EarlyStopping(monitor='val_acc', min_delta=0.001,
                               patience=patience, verbose=0, mode='max')
#Build model
Model=ModelBuilder(DWIshape, Mapshape, dropout=self.Dropout,
                      regularization=self.Regularization,
                      activ='relu', DWIconv=self.nDWI, DWIsize=self.sDWI,
                      classes=classnum, layers=self.nCNN,
                      filtersize=self.sCNN,
                      FClayers=self.FCL, last=self.Last)
#Compile
Model.compile(optimizer=K.optimizers.Adam(lr,decay=self.Decay), loss=loss, metrics=['accuracy'])
#Fit
his=Model.fit(x=TRAINDATA,y=TrueTrain,epochs=epochs,batch_size=5, shuffle=True, validation_data=(TESTDATA,TrueTest), verbose=0, callbacks=[STOP, MemoryCallback()]) #check verbose and callbacks
#Extract 
S=Model.evaluate(x=TESTDATA, y=TrueTest,verbose=1)[1]
del his
del Model
del rateme
K.backend.clear_session()
tf.reset_default_graph()
gc.collect()

return S

Answer 1

由于使用model.fit()等内置函数时，TensorFlow 2.4.1中似乎仍然存在内存泄漏，这是我的看法。

问题

• 尽管我运行的是 NVIDIA GeForce RTX 2080 TI GPU，但 RAM 使用量仍然很大。
• 随着训练的进行，增加纪元时间。
• 某种内存泄漏（感觉有点线性）。

解决方案

• 将run_eagerly=True参数添加到函数中model.compile()。但是，这样做可能会导致 TensorFlow 的图优化不再起作用，从而导致性能下降（参考）。
• 创建一个自定义回调，在每个纪元结束时进行垃圾收集并清除 Keras 后端（参考）。
• 不要activation使用tf.keras.layers. 将激活函数作为单独的层（参考）。
• 使用LeakyReLU而不是ReLU作为激活函数（参考）。

注意：由于所有要点都可以单独实现，因此您可以混合和匹配它们，直到获得适合您的结果。不管怎样，这里有一个代码片段，展示了所有的解决方案：

import gc
from tensorflow.keras import backend as k
from tensorflow.keras.layers import Conv2D, BatchNormalization, ReLU
from tensorflow.keras.callbacks import Callback


class CovNet:
    ...
    x = Conv2d(
        ...,
        activation=None
    )(x)
    x = BatchNormalization()(x)
    x = ReLU()(x)  # or LeakyReLU
    ...

#--------------------------------------------------------------------------------

class ClearMemory(Callback):
    def on_epoch_end(self, epoch, logs=None):
        gc.collect()
        k.clear_session()

#--------------------------------------------------------------------------------

model.compile(
    ...,
    run_eagerly=True
)

#--------------------------------------------------------------------------------

model.fit(
    ...,
    callbacks=ClearMemory()
)

通过这些解决方案，我现在能够在占用更少 RAM 的情况下进行训练，纪元时间保持不变，如果仍然存在内存泄漏，则可以忽略不计。

感谢@Hongtao Yang 提供了相关GitHub 问题之一的链接，并感谢 GitHub 上的 rschiewer 发表评论。

笔记

• 如果以上方法都不适合您，您可能需要尝试在 TensorFlow 中编写自己的训练循环。这是有关如何操作的指南。
• 人们还报告说，使用tcmalloc而不是默认malloc分配器在一定程度上减轻了内存泄漏。如需参考，请参阅此处或此处。

我希望这对其他人也有帮助，并为您节省一些在互联网上进行研究的时间。