TF2 将 report_tensor_allocations_upon_oom 添加到 RunOptions

Question

我收到这条消息：

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 中执行此操作的工作示例。

Answer 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