如何测量 TensorFlow 模型的 GPU 内存使用情况答案

【问题标题】：How to measure GPU memory usage of TensorFlow model如何测量 TensorFlow 模型的 GPU 内存使用情况
【发布时间】：2020-01-07 15:56:38
【问题描述】：

我的目标是弄清楚在推理过程中保存为 .pb 文件的 TensorFlow 模型使用了多少 GPU 内存。通过将per_process_gpu_memory_fraction 的值限制为 0.01 并在输入图像上成功运行模型，我预计内存使用量为 120MB（基于 12,000MB GPU）。

conf = tf.ConfigProto()
conf.gpu_options.per_process_gpu_memory_fraction=0.01
session = tf.Session(config=conf)

当我通过nvidia-smi 观察实际 GPU 内存使用情况时，我看到使用率为 550MB。根据此处的答案 (https://stackoverflow.com/a/55532954)，我测量了创建 TF 会话所需的默认内存，即 150MB。现在的问题是剩余的 550MB（测量）- 150MB（会话）- 120MB（模型）= 280MB 来自哪里。我想排除使用 GPU 内存的任何其他因素，只量化模型本身使用了多少。在这种情况下我可以简单地假设 120MB 吗？

【问题讨论】：

标签： python tensorflow gpu

【解决方案1】：

在运行 python 代码之前尝试将 TF_FORCE_GPU_ALLOW_GROWTH 设置为 true：

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

在你导入 tensorflow 之前放这个。

然后测量 GPU 内存利用率。此外，它还取决于 Tensorflow 内存分配器如何分配 GPU 上的内存。所以它通常会比内存中的实际模型大一点。

也可以试试这里提到的建议：How to restrict tensorflow GPU memory usage?

【讨论】：

我尝试了您的建议，但通过 nvidia-smi 显示的值与默认内存使用和预期模型内存使用的总和之间的差异仍然存在。
现在的内存消耗是多少？
和以前一样。我认为这是因为选择了 per_process_gpu_memory_fraction 的值，所以它正好适合模型。