TF2 将 report_tensor_allocations_upon_oom 添加到 RunOptions答案

【问题标题】：TF2 add report_tensor_allocations_upon_oom to RunOptionsTF2 将 report_tensor_allocations_upon_oom 添加到 RunOptions
【发布时间】：2021-01-19 15:24:58
【问题描述】：

我收到了这条消息：

Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

如何在 TensorFlow 2.3 中做到这一点？

在过去的几天里，这被证明是一个令人惊讶的令人沮丧的问题。似乎没有关于如何在 TF2 中执行此操作的有效示例。

【问题讨论】：

这有帮助吗？ stackoverflow.com/questions/49665757/…
@rvinas 我很遗憾目前无法测试它，但我记得看过那个链接。它没有为我解决问题。 Iirc 因为它是针对 TF1 的，并且不兼容或转换为 TF2。很想被证明是错误的。
有关于这个问题的消息吗？使用 tf1.compat 时，我仍然使用 TypeError: Invalid keyword argument(s) in 'compile': {'run_metadata', 'options'}，我不知道如何处理其他建议的答案...

标签： tensorflow tensorflow2.0

【解决方案1】：

这距离分配的张量列表还有很长的路要走，但对于 TF2 来说是一个开始：

Tensorflow 2.4.1 包含 tf.config.experimental.get_memory_usage 方法，该方法返回 GPU 上当前使用的字节数。比较不同时间点的这个值可以揭示哪些张量占用了 VRAM。看起来很准确。

顺便说一句，最新的夜间版本包含 tf.config.experimental.get_memory_info 方法，似乎他们改变了主意。这个包含current，以及peak 使用的内存。

TF 2.4.1 上的示例代码：

import tensorflow as tf

print(tf.config.experimental.get_memory_usage("GPU:0"))  # 0

tensor_1_mb = tf.zeros((1, 1024, 256), dtype=tf.float32)
print(tf.config.experimental.get_memory_usage("GPU:0"))  # 1050112

tensor_2_mb = tf.zeros((2, 1024, 256), dtype=tf.float32)
print(tf.config.experimental.get_memory_usage("GPU:0"))  # 3147264

tensor_1_mb = None
print(tf.config.experimental.get_memory_usage("GPU:0"))  # 2098688

tensor_2_mb = None
print(tf.config.experimental.get_memory_usage("GPU:0"))  # 1536

【讨论】：

【解决方案2】：

为了加载 RunOptions，您需要使用 RunMetadata。 TF2 的这两个文件都可以在 tf.compat.v1 包中找到。以下代码适用于带有 tensorflow 2.4.0 后端的 keras 2.4.3：

#Model build()
import tensorflow as tf
run_opts = tf.compat.v1.RunOptions(report_tensor_allocations_upon_oom=True)
runmeta = tf.compat.v1.RunMetadata()
keras_model.compile(optimizer=..., loss=..., options = run_opts, run_metadata=runmeta)
#Model fit()

【讨论】：

这导致TypeError: Invalid keyword argument(s) in 'compile': {'run_metadata', 'options'}