识别 tensorflow 数据管道和训练中的内存泄漏？答案

【问题标题】：Identify memory leak in tensorflow data pipeline and training?识别 tensorflow 数据管道和训练中的内存泄漏？
【发布时间】：2019-12-07 04:38:24
【问题描述】：

TF.Data.Dataset 管道的内存泄漏。是否有分析器来识别管道或 tf.keras 训练中的内存泄漏？几个问题，如果你有任何想法—— 1.我忽略的伪代码有明显问题吗？ 2. 关于在哪里/寻找什么有什么想法吗？ 3. 任何关于如何在训练继续精确定位问题时对内存使用情况进行分析的指针？

我刚刚在 Tensorflow 1.15 下将我的代码库转移到了 Eager 模式，但我遇到了以前没有的内存问题。在进入渴望模式之前，我可以毫无问题地训练 500 多个 epoc，而现在，训练在 70 个 epoc 后停止。我正在尝试找出一种方法来确定泄漏的位置，我希望你们中的一些人有一些想法。

我正在使用 tf.data.Dataset 来构建数据管道（请参见下面的伪代码）并加快数据馈送，我正在使用具有交错的数据集，如下所示。我已经对存储在分片 TFRecords 文件中的数据进行了预处理，并且数据集 API 会加载数据并进行最少的处理以提供适当的批量大小的数据。 GPU 内存似乎很好，并且训练一直持续到 CPU RAM 完全耗尽。如下表所示，psutil 内存日志显示 CPU RAM 持续增加。

我尝试过的：按照这些建议明确调用 gc.collect, tf.set_random_seed(1) 但似乎没有帮助。 https://github.com/tensorflow/tensorflow/issues/30324 Memory Continually Increasing When Iterating in Tensorflow Eager Execution

Ubuntu 18.04，tf-nightly-gpu 1.15.0.dev20190728 CPU - AMD Ryzen Threadripper 1920X 12 核处理器内存 – 128GB GPU - RTX 2080 Ti 11GB

#Generator that is passed to the fit_generator
    def get_simple_Dataset_generator(….):

    dataset = load_dataset (…)
          while True:
               try:
                 for x, Y in dataset:
                   yield x, Y
               finally:
                     dataset = load_dataset (“change data sources”)
                     #tried gc.collect(), tf.set_random_seed(1)

#sets up the dataset with interleave. 
def load_dataset(…):
   #setup etc
   dataset = dataset.interleave(lambda x: 
   tf.data.Dataset.from_generator(self.simple_gen_step1, 
                                output_types=(tf.string, tf.float32, tf.string), args=(x,batch_size, lstm_reshape,)),
                                cycle_length=2,
                                block_length=1)

    dataset = dataset.interleave(lambda each_ticker, each_dataset, each_dates: tf.data.Dataset.from_generator(self.simple_gen_step2, 
                                output_types=(tf.float32, tf.int16), args=(names, dataset, dates,batch_size,)),
                                cycle_length=2,
                                block_length=1)

  return dataset

#Our Model uses CuDNNLSTM and Dense layers
build_model():
    model = Sequential()
    model.add(CuDNNLSTM(feature_count,
                            batch_input_shape=(batch_size,look_back, feature_count),
                            stateful=Settings.get_config(Settings.STATEFUL), 
                            return_sequences=True))

    model.add(CuDNNLSTM(feature_count, return_sequences = True))
    model.add(CuDNNLSTM(feature_count, return_sequences = True))
    model.add(CuDNNLSTM(feature_count, return_sequences = False))
    model.add(Dropout(dropout))

    model.add(Dense( max(feature_count//(2*1), target_classes), use_bias=False))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(dropout))

    model.add(Dense(max(feature_count//(2*2), target_classes),use_bias=False))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(dropout))

    model.add(Dense(max(feature_count//(2*3), target_classes), use_bias=False))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(dropout))

    model.add(Dense(target_classes, activation='softmax'))

    return model

psutil 日志中显示的 CPU RAM

【问题讨论】：

标签： tensorflow memory-leaks tensorflow-datasets

【解决方案1】：

对于遇到类似问题的任何人，我认为如果在 fit_generator 中使用 class_weights，则会出现内存泄漏。我已经发布了另一个更详细的信息。 Using class_weights in fit_generator causes memory leak

【讨论】：

【解决方案2】：

我想我会分享我发现的关于 TensorFlow 2.x.x 中的内存泄漏。它可能不是对您的具体问题的 100 % 具体答案，但它可能会帮助其他人在使用 model.fit() 等内置函数时解决他们的内存泄漏问题。

这是相关GitHub issues 之一的链接，这是我的solution（请同时考虑我的解决方案的 cmets）。

【讨论】：