Keras在每个历元中占据无限增加的内存量



我正在运行一个遗传超参数搜索算法,它会迅速饱和所有可用内存。

经过几次测试后,keras所需的内存量似乎在不同时期之间以及训练不同模型时都有所增加。随着小批量的增加,问题变得更加严重,1~5的小批量至少给了我足够的时间,让我看到内存使用量在最初几次快速增长,然后随着时间的推移缓慢但稳定地增加。

我已经检查了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

由于内存泄漏在使用model.fit((等内置函数时,TensorFlow 2.4.1中似乎仍然存在


问题

  • 即使我运行的是NVIDIA GeForce RTX 2080 TI GPU,RAM的使用量也很高
  • 随着训练的进行,历元次数不断增加
  • 某种内存泄漏(感觉有点线性(

解决方案

  • run_eagerly=True参数添加到model.compile()函数中。然而,这样做可能会导致TensorFlow的图优化不再起作用,这可能会导致性能下降(参考(
  • 创建一个自定义回调,在每个epoch(引用(结束时垃圾收集并清除Keras后端
  • 不要在tf.keras.layers中使用activation参数。将激活函数作为一个单独的层(参考(
  • 使用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分配器在一定程度上减轻了内存泄漏。有关参考资料,请参阅此处或此处

我希望这也能帮助其他人,并为你节省一些在互联网上研究的时间。

消耗整个可用内存是TF的默认行为。
您可以使用以下代码限制TF中的内存消耗量:

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.9 # fraction of memory
config.gpu_options.visible_device_list = "0"
set_session(tf.Session(config=config))

最后,我只是用bash脚本在每个训练会话之间重新启动了python会话,找不到更好的方法来避免内存占用率爆炸的

也许这是一个相关的问题?如果是这样,那么使用自定义训练循环而不是model.fit方法时就可以了。

我认为他们还没有解决这个问题,所以我会避免使用内置的训练/评估/预测方法。

最新更新