如何在 TensorFlow 分析器中获取详细的内存细分？

Posted 2023-03-06

【中文标题】如何在 TensorFlow 分析器中获取详细的内存细分？

【英文标题】：How to get detailed memory breakdown in the TensorFlow profiler?

【发布时间】：2017-10-26 18:11:31

【问题描述】：

我正在使用新的TensorFlow profiler 来分析我的神经网络中的内存使用情况，我在具有 12GB RAM 的 Titan X GPU 上运行。这是我分析主要训练循环时的一些示例输出：

==================Model Analysis Report======================
node name | requested bytes | ...
Conv2DBackpropInput        10227.69MB (100.00%, 35.34%),     ...
Conv2D                       9679.95MB (64.66%, 33.45%),     ...
Conv2DBackpropFilter         8073.89MB (31.21%, 27.90%),     ...

显然，这加起来超过 12GB，因此其中一些矩阵必须在主内存中，而其他矩阵则在 GPU 上。我很想看到在给定步骤中哪些变量在哪里的详细分类。是否可以通过分析器或其他方式获得有关各种参数存储位置（主内存或 GPU 内存）的更详细信息？

【参考方案1】：

“请求的字节数”显示所有内存分配的总和，但可以分配和取消分配内存。所以仅仅因为“请求的字节数”超过了 GPU RAM 并不一定意味着内存正在被传输到 CPU。

特别是对于前馈神经网络，TF 通常会保留前向激活，以提高反向传播效率，但不需要保留中间反向传播激活，即每一层的 dL/dh，因此它可以只完成这些中间体后丢弃这些中间体。所以我认为在这种情况下你关心的是 Conv2D 使用的内存，它小于 12 GB。

您还可以使用时间线来验证总内存使用量从未超过 12 GB。