如何在 tensorflow2-gpu 中调试内存不足

Question

我正在使用 tensorflow-2 gpu 和 tf.data.Dataset。 对小型模型进行训练。

当训练一个更大的模型时，一开始一切正常：使用 gpu，第一个 epoch 工作没有问题（但我正在使用我的大部分 gpu 内存）。

在验证时，我遇到了一个 CUDA_ERROR_OUT_OF_MEMORY，它有各种分配，无法分配的字节数越来越少（范围从 922Mb 到 337Mb）。 我目前没有指标，也没有回调，我正在使用tf.keras.Model.fit。 如果我删除验证数据，训练会继续。

我的问题是什么？我该如何调试？

在 tf1 中，我可以使用 RunOptions(report_tensor_allocations_upon_oom=True)，在 tf2 中是否存在任何等价物？

这发生在 tensorflow==2.1.0 中。

Answer 1

这些没有发生在 2.0 alpha TensorFlow 但在 2.0.

Pip 安装 tensorflow-gpu==2.0.0: 有 内存泄露！
pip install tensorflow-gpu==2.0.0-alpha： 没关系！

试试看