如何确保训练阶段不会面临 OOM？

Question

标题中的问题已完成。“如何确保训练阶段不会面临 OOM？”

只是一些旁注，根据我的经验，有两种 OOM 情况。一种是当您的模型和小批量所需的内存大于您拥有的内存时。在这种情况下，训练阶段永远不会开始。解决此问题的解决方案是使用较小的批量。尽管如果我能计算出我的硬件可以为某个特定模型管理的最大批量大小，那就太好了。但即使我在第一次尝试时找不到最大的批量大小，我也总能通过反复试验找到它（因为该过程立即失败）。

我面临 OOM 的第二种情况是训练过程开始时，它会持续一段时间。甚至可能是几个时代。但后来不知什么原因，它面临着OOM。对我来说，这种情况令人沮丧。因为它随时可能发生，您永远不会知道正在进行的培训是否会结束。到目前为止，我已经失去了几天的训练时间，而我认为一切都在进行中。

我认为一些澄清是有序的。首先，我说的是带有 GPU 的个人计算机。其次，GPU专用于计算，不用于显示。如果我错了，请纠正我，但我相信这意味着训练过程在不同的时间点需要不同的内存大小。怎么会这样？再一次，我如何确保我的训练阶段不会面临 OOM？

以这次跑步为例：

3150/4073 [======================>.......] - ETA: 53:39 - loss: 0.3323
2019-10-13 21:41:13.096320: W tensorflow/core/common_runtime/bfc_allocator.cc:314] Allocator (GPU_0_bfc) ran out of memory trying to allocate 60.81MiB (rounded to 63766528).  Current allocation summary follows.

经过三个小时的训练，TensorFlow 要求的内存超出了我的硬件所能提供的范围。我的问题是，为什么在这一点而不是在进程开始时增加内存分配？

[更新]

鉴于 Eager 模式的已知问题，我将对我的案例进行一些说明。我不是在急切模式下编码。这是我的训练代码的样子：

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
training_dataset = tf.data.Dataset.from_tensor_slices(...)
validation_dataset = tf.data.Dataset.from_tensor_slices(...)

with strategy.scope():
    model = create_model()

    model.compile(optimizer='adam', loss='categorical_crossentropy')

    pocket = EarlyStopping(monitor='val_loss', min_delta=0.001,
                           patience=5, verbose=1,
                           restore_best_weights = True)

    history = model.fit(training_dataset.shuffle(buffer_size=1000).batch(30),
                        epochs=3,
                        callbacks=[pocket],
                        validation_data=validation_dataset.shuffle(buffer_size=1000).batch(30),
                        workers=3, use_multiprocessing=True)

Answer 1

如果您在循环中重复训练，就会发生已知的内存泄漏 [1]。解决这个问题的方法是调用tf.keras.backend.clear_session()，并且可能gc.collect()时不时地在循环中调用。

不过，TF 1.15 和 2.0 中的行为似乎有所不同，这可能不足以修复它。我发现tf.keras.backend.clear_session()在 CPU 上的训练循环中会重置逐渐的内存泄漏，而不会影响训练。

[1] https://github.com/tensorflow/tensorflow/issues/30324