火车+验证循环上的tf.Estimator RAM（MemoryError）溢出

Question

我有一个非常标准的tf.Estimator代码，带有训练/评估循环，如下所示：

for _ in range(params['epochs']):
    classifier.train(input_fn=lambda: input_fn(example_generator, source='train', batch_size=params['batch_size']))
    classifier.evaluate(input_fn=lambda: input_fn(example_generator, source='validation', batch_size=params['batch_size']))

首先要占用大约5GB的RAM（包括所有其他CPU任务），并且在不加干扰的情况下填满整个16GB可用空间。 模型正在GPU上训练，但是CPU溢出了。

我不明白为什么运行此代码时RAM使用率应该改变。每进行.evaluate / .train个通话，就会再累积200MB〜。

如果有什么不同，input_fn如下所示：

def input_fn(example_generator: ExampleGenerator, source, batch_size):
    dataset = tf.data.Dataset.from_generator(generator=lambda: example_generator.generate(source, shuffle=True),
                                             output_types=(tf.string, tf.int32, tf.int32, tf.int32, tf.int32),
                                             output_shapes=(
                                                 tf.TensorShape([]), tf.TensorShape([None]), tf.TensorShape([None]),
                                                 tf.TensorShape([None]), tf.TensorShape([])))

    # Shuffle, repeat, and batch the examples.
    dataset = dataset.padded_batch(batch_size,
                                                 padded_shapes=([], [None], [None], [None], []))

    dataset = dataset.prefetch(buffer_size=batch_size)

    iterator = dataset.make_one_shot_iterator()

    f1, f2, f3, f4, label = iterator.get_next()

    return {'f1': f1, 'f2': f2, 'f3': f3, 'f4': f4}, label

ExampleGenerator是一个仅包含所有示例列表的简单类，其中的generate方法可产生示例并在两次调用之间将它们随机播放。

编辑：

我还尝试过将ExampleGenerator类与生成器函数交换，该函数产生相同的输出，相同的结果。