如何将 report_tensor_allocations_upon_oom 添加到 Keras 中的 RunOptions

Question

我正在尝试使用 Keras 在 GPU 上训练神经网络，但出现“资源耗尽：分配张量时出现 OOM”错误。它试图分配的特定张量不是很大，所以我假设之前的一些张量几乎消耗了所有 VRAM。错误消息带有一个暗示：

提示：如果您想在 OOM 发生时查看已分配张量的列表，请将 report_tensor_allocations_upon_oom 添加到 RunOptions 以获取当前分配信息。

这听起来不错，但我该怎么做呢？ RunOptions 似乎是 Tensorflow 的东西，我能找到的关于它的少量文档将它与“会话”相关联。我使用的是 Keras，因此 Tensorflow 隐藏在一个抽象层之下，而它的会话则隐藏在该抽象层之下的另一层之下。

我如何挖掘所有内容以设置此选项使其生效？

原文由 dspeyer 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

TF1解决方案：

它并不像看起来那么难，您需要知道的是，根据 文档，传递给 model.compile 的 **kwargs 参数将传递给 session.run

所以你可以这样做：

 import tensorflow as tf
run_opts = tf.RunOptions(report_tensor_allocations_upon_oom = True)

model.compile(loss = "...", optimizer = "...", metrics = "..", options = run_opts)

每次调用 session.run 时应该直接传递。

TF2：

上面的解决方案仅适用于 tf1。不幸的是，对于 tf2，似乎还没有简单的解决方案。

原文由 Dr. Snoopy 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

目前，无法将选项添加到 model.compile 。参见： https ://github.com/tensorflow/tensorflow/issues/19911

原文由 Richard 发布，翻译遵循 CC BY-SA 4.0 许可协议