tensorflow - 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