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

时间:2018-12-05 19:33:44

标签: python tensorflow tensorflow-estimator

我有一个非常标准的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类与生成器函数交换,该函数产生相同的输出,相同的结果。

0 个答案:

没有答案