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

Question

我使用的是带有tf.data.Dataset的tensorflow-2图形处理器。在小模型上进行训练是可行的。

当训练一个更大的模型时，首先，一切都可以正常工作:使用gpu，第一个时期工作没有问题(但我使用了大部分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:没问题！

试试看