TensorFlow：如何测量每个张量占用多少 GPU 内存？

Posted 2023-02-23

【中文标题】TensorFlow：如何测量每个张量占用多少 GPU 内存？

【英文标题】：TensorFlow: How to measure how much GPU memory each tensor takes?

【发布时间】：2016-07-19 18:57:25

【问题描述】：

我目前正在 TensorFlow 中实现 YOLO，我对占用的内存量感到有些惊讶。在我的 GPU 上，我可以使用他们的 Darknet 框架以批量大小 64 训练 YOLO。在 TensorFlow 上，我只能以批量大小 6 来训练，而 8 的内存已经用完。对于测试阶段，我可以使用 64 的批处理大小运行而不会耗尽内存。

我想知道如何计算每个张量消耗了多少内存？默认情况下所有张量都保存在 GPU 中吗？我可以简单地将总内存消耗计算为形状 * 32 位吗？

我注意到，由于我使用了动量，所以我所有的张量也有一个 /Momentum 张量。会不会也占用了很多内存？

我正在使用distorted_inputs 方法扩充我的数据集，该方法与CIFAR-10 tutorial 中定义的方法非常相似。莫非这部分占用了很大的内存？我相信 Darknet 会在 CPU 中进行修改。

【参考方案1】：

现在1258 已关闭，您可以在导入 TensorFlow 之前通过设置环境变量在 Python 中启用内存日志记录：

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

因此会有很多的日志记录。您需要对结果进行 grep 查找以找到适当的行。例如：

grep MemoryLogTensorAllocation train.log

【讨论】：

这是在 python 中，但我怎样才能抑制所有登录 tensorflow cpp

【参考方案2】：

抱歉回复缓慢。不幸的是，现在设置日志级别的唯一方法是编辑 tensorflow/core/platform/logging.h 并使用例如重新编译

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

1258 有一个 bug 可以更优雅地控制日志记录。

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

【参考方案3】：

请参阅此 (commit) 中的说明。 内存分配是原始信息，尽管它需要一个脚本以易于阅读的形式收集信息。

【讨论】：

感谢您的快速回复。我对如何使用 python 设置 VLOG 有点困惑？能否请你帮忙？谢谢！

什么是MemoryLogTensorOutput？