TensorFlow:如何测量每个张量所需的GPU内存？

Question

我目前正在用YOLO实现TensorFlow，对于占用多少内存我感到有点惊讶。在我的GPU上，我可以使用批处理大小为64的暗网框架来训练YOLO。在TensorFlow上，我只能在批处理大小为6的情况下才能这样做，而我已经用光了内存。在测试阶段，我可以在不耗尽内存的情况下以批处理大小64运行。

1. 我想知道如何计算每个张量消耗了多少内存？默认情况下，所有张量都保存在GPU中吗？我可以简单地计算总的内存消耗的形状* 32位吗？
2. 我注意到，由于我使用动量，我所有的张量都有一个/Momentum张量。这也会占用大量的内存吗？
3. 我使用一个方法distorted_inputs来增强我的数据集，非常类似于在CIFAR-10教程中定义的方法。是不是这部分占用了很大一部分记忆？我相信Darknet会在CPU中做修改。

Answer 1

现在一二五八已经关闭，您可以通过在导入TensorFlow之前设置一个环境变量来启用TensorFlow中的内存日志记录：

import os
os.environ['TF_CPP_MIN_VLOG_LEVEL']='3'
import tensorflow as tf

这将导致大量的日志记录。您将希望grep的结果找到适当的行。例如：

grep MemoryLogTensorAllocation train.log

Answer 2

很抱歉回答得太慢了。不幸的是，现在设置日志级别的唯一方法是编辑tensorflow/core/platform/logging.h，然后用.

#define VLOG_IS_ON(lvl) ((lvl) <= 1)

有一个打开一二五八的bug，可以更优雅地控制日志记录。

MemoryLogTensorOutput条目在每次Op执行结束时都会被记录下来，并指示保存Op输出的张量。知道这些张量是有用的，因为内存直到下游Op消耗掉这些张量才会释放，这在一个大图中可能要晚得多。

Answer 3

请参阅本文(提交)中的描述。内存分配是原始信息，虽然它需要一个脚本来收集信息，以一个易于阅读的形式。